test_file_reader.py

#  Copyright (C) 2022 ASTRON (Netherlands Institute for Radio Astronomy)
#  SPDX-License-Identifier: Apache-2.0

from os.path import dirname
from typing import List, Dict

from numpy import ndarray

from lofar_station_client.file_access import member, read_hdf5, attribute
from tests import base


class DataSubSet(object):
    values: List[int] = member()


class DataSet(object):
    observation_station: str = attribute()
    observation_source: str = attribute(from_member="sub_set")
    nof_payload_errors: List[int] = member()
    nof_valid_payloads: List[int] = member()
    values: List[List[float]] = member()
    sub_set: DataSubSet = member(name="test")
    non_existent: DataSubSet = member(optional=True)

    def __repr__(self):
        return f"DataSet(nof_payload_errors={self.nof_payload_errors})"


class AttrDataSet(object):
    observation_station: str = attribute()
    observation_station_optional: str = attribute(optional=True)
    observation_station_missing_none_optional: str = attribute(optional=False)
    test_attr: str = attribute(from_member="calibration_data", name="test_attribute")
    calibration_data: ndarray = member(name="data")


class CalData:
    x_attr: str = attribute("test_attr", from_member="x")
    y_attr: str = attribute("test_attr", from_member="y")
    x: ndarray = member()
    y: ndarray = member()


class CalTable(Dict[str, CalData]):
    observation_station: str = attribute()


class CalTableDict(Dict[str, Dict[str, ndarray]]):
    pass


class TestHdf5FileReader(base.TestCase):
    def test_file_reading(self):
        with read_hdf5(
            dirname(__file__) + "/SST_2022-11-15-14-21-39.h5", Dict[str, DataSet]
        ) as ds:
            self.assertEqual(21, len(ds.keys()))
            item = ds["SST_2022-11-15T14:21:59.000+00:00"]
            self.assertEqual(
                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                item.nof_payload_errors,
            )
            # double read to check if (cached) value is the same
            self.assertEqual(
                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                item.nof_payload_errors,
            )
            self.assertEqual(
                [12, 12, 12, 12, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                item.nof_valid_payloads,
            )
            self.assertIsNone(item.non_existent)
            self.assertEqual(192, len(item.values))

            with self.assertRaises(KeyError):
                _ = (
                    item.sub_set
                    # this item does not exist but is not marked as optional
                )

            item = ds["SST_2022-11-15T14:21:39.000+00:00"]
            self.assertEqual(100, len(item.sub_set.values))

    def test_read_attribute(self):
        with read_hdf5(dirname(__file__) + "/cal-test.h5", AttrDataSet) as ds:
            self.assertEqual("test-station", ds.observation_station)
            self.assertEqual("dset_attr", ds.test_attr)
            self.assertIsNone(ds.observation_station_optional)
            self.assertEqual("dset_attr", ds.test_attr)  # test caching
            with self.assertRaises(KeyError):
                _ = (
                    ds.observation_station_missing_none_optional
                    # this attribute does not exist but is not marked as optional
                )

    def test_read_ndarray(self):
        with read_hdf5(dirname(__file__) + "/cal-test.h5", AttrDataSet) as ds:
            d = ds.calibration_data
            self.assertTrue(isinstance(d, ndarray))
            self.assertEqual(512, d.shape[0])
            self.assertEqual(96, d.shape[1])

    def test_read_derived_dict(self):
        with read_hdf5(dirname(__file__) + "/cal-test-dict.h5", CalTable) as ds:
            self.assertEqual(5, len(ds))
            self.assertEqual("test-station", ds.observation_station)
            ant_2 = ds["ant_2"]
            self.assertEqual(512, len(ant_2.x))
            self.assertEqual(512, len(ant_2.y))
            self.assertEqual("ant_2_x", ant_2.x_attr)
            self.assertEqual("ant_2_y", ant_2.y_attr)

    def test_read_derived_double_dict(self):
        with read_hdf5(dirname(__file__) + "/cal-test-dict.h5", CalTableDict) as ds:
            self.assertEqual(5, len(ds))
            ant_2 = ds["ant_2"]
            self.assertTrue("x" in ant_2)
            self.assertTrue("y" in ant_2)
            self.assertEqual(512, len(ant_2["x"]))
            self.assertEqual(512, len(ant_2["y"]))

    def test_read_as_object(self):
        class ObjectDataSet:
            item_1: Dict[str, List[int]] = member(
                name="SST_2022-11-15T14:21:59.000+00:00"
            )
            item_2: Dict[str, List[int]] = member(
                name="SST_2022-11-15T14:21:39.000+00:00"
            )

        with read_hdf5(
            dirname(__file__) + "/SST_2022-11-15-14-21-39.h5", ObjectDataSet
        ) as ds:
            self.assertEqual(
                ["nof_payload_errors", "nof_valid_payloads", "values"],
                list(ds.item_1.keys()),
            )
            with self.assertRaises(TypeError):
                _ = (
                    ds.item_2["test"]
                    # item test is of type group and will raise an error
                )

    def test_malformed_data(self):
        class BrokenDataSet(object):
            nof_payload_errors: DataSubSet = member()
            nof_valid_payloads: int = member()
            sub_set: List[int] = member(name="test")

        with read_hdf5(
            dirname(__file__) + "/SST_2022-11-15-14-21-39.h5", Dict[str, BrokenDataSet]
        ) as ds:
            item = ds["SST_2022-11-15T14:21:39.000+00:00"]
            with self.assertRaises(TypeError):
                _ = item.nof_payload_errors
            with self.assertRaises(TypeError):
                _ = item.nof_valid_payloads
            with self.assertRaises(TypeError):
                _ = item.sub_set

    def test_reader_close(self):
        file_reader = read_hdf5(
            dirname(__file__) + "/SST_2022-11-15-14-21-39.h5", Dict[str, DataSet]
        )
        file_reader.close()