Fix tests and make it easier to add more opal data

datahub/main.py

@@ -5,7 +5,7 @@
 from . import data as dt
 from . import log
 from .dsr import dsr_headers, read_dsr_file, validate_dsr_data
-from .opal import OpalArrayData, OpalModel
+from .opal import OpalArrayData, OpalModel, opal_headers
 from .wesim import get_wesim
 
 app = FastAPI(
@@ -35,15 +35,19 @@ def create_opal_data(data: OpalModel | OpalArrayData) -> dict[str, str]:
     if isinstance(data, OpalArrayData):
         log.info("Array format detected.")
         append_input = raw_data["array"]
+        del append_input[5:8]
+
+        if len(append_input) != len(opal_headers) + 1:
+            message = (
+                f"Array has invalid length. Expecting {len(opal_headers) + 4} items."
+            )
+            log.error(message)
+            raise HTTPException(status_code=400, detail=message)
+
     else:
         log.info("Dict format detected.")
         append_input = raw_data
 
-    if isinstance(append_input, list) and not len(append_input) == 45:
-        message = "Array has invalid length. Expecting 45 items."
-        log.error(message)
-        raise HTTPException(status_code=400, detail=message)
-
     log.info("Appending new data...")
     log.debug(f"Original Opal DataFrame:\n\n{dt.opal_df}")
     try:

datahub/opal.py

@@ -66,7 +66,7 @@ class OpalModel(BaseModel):
     bm_cost_less: float = Field(alias="Balancing Market Cost (Less)")
     bm_cost_least: float = Field(alias="Balancing Market Cost (Least)")
     bm_power_most: float = Field(alias="Balancing Market Power (Most)")
-    bm_power_more: float = Field(alias="Balancing Market Power (Most)")
+    bm_power_more: float = Field(alias="Balancing Market Power (More)")
     bm_power_current: float = Field(alias="Balancing Market Power (Current)")
     bm_power_less: float = Field(alias="Balancing Market Power (Less)")
     bm_power_least: float = Field(alias="Balancing Market Power (Least)")
@@ -77,7 +77,7 @@ class Config:
         allow_population_by_field_name = True
 
 
-opal_headers = {
+opal_headers: dict[str, str] = {
     field["title"]: name
     for name, field in OpalModel.schema(by_alias=False)["properties"].items()
     if name != "frame"
@@ -119,7 +119,11 @@ def append(self, data: dict[str, int | float] | list[int | float]) -> None:
             data_index = data[0]
         else:
             data_index = data["frame"]
+
+        dtypes = self._obj.dtypes
         self._obj.loc[data_index] = row.loc[data_index]  # type: ignore[call-overload]
+        self._obj[:] = self._obj.astype(dtypes)[:]
+        self._obj[self._obj.columns] = self._obj.astype(dtypes)[self._obj.columns]
 
 
 def create_opal_frame() -> pd.DataFrame:
@@ -128,8 +132,15 @@ def create_opal_frame() -> pd.DataFrame:
     Returns:
         An initial Dataframe for the opal data with key frame 0
     """
-    df = pd.DataFrame(0, index=range(1), columns=list(opal_headers.keys()))
-    df["Time"] = pd.Timestamp(OPAL_START_DATE).as_unit("ns")  # type: ignore[attr-defined]  # noqa: E501
+    dtypes = {
+        field["title"]: ("int" if field["type"] == "integer" else "float")
+        for name, field in OpalModel.schema(by_alias=False)["properties"].items()
+        if name != "frame"
+    }
+    dtypes["Time"] = "datetime64[ns]"
+
+    df = pd.DataFrame(0, index=range(0), columns=list(opal_headers.keys()))
+    df = df.astype(dtypes)
 
     return df
 
@@ -149,10 +160,7 @@ def get_opal_row(data: dict[str, int | float] | list[int | float]) -> pd.DataFra
 
     else:
         data_array = data.copy()
-        data_index = data_array[0]
-
-        del data_array[5:8]
-        del data_array[0]
+        data_index = data_array.pop(0)
 
     row = pd.DataFrame(
         [data_array], index=[data_index], columns=list(opal_headers.keys())

tests/conftest.py

@@ -24,15 +24,14 @@ def opal_data():
     data["frame"] = 1
     for key in list(opal_headers.values()):
         data[key] = np.random.randint(100)
+    data["time"] = 8.58
     return data
 
 
 @pytest.fixture
-def opal_data_array():
+def opal_data_array(opal_data):
     """Pytest Fixture for random Opal data input in array format."""
-    data = [1, 8.58]
-    for x in range(43):
-        data.append(np.random.randint(100))
+    data = list(opal_data.values())
     return data
 
 

tests/test_opal.py

@@ -1,31 +1,28 @@
+import numpy as np
 import pandas as pd
 import pytest
 
 
 def test_create_opal_frame():
     """Tests creation of blank Opal Dataframes."""
-    from datahub.opal import create_opal_frame
+    from datahub.opal import create_opal_frame, opal_headers
 
     df = create_opal_frame()
 
     # Checks that Dataframe has 41 columns and only 1 row.
-    assert len(df.columns) == 41
-    assert len(df.index) == 1
-
-    time_data = df.iloc[0]["Time"]
+    assert df.shape == (0, len(opal_headers))
 
-    # Checks that 'Time' data is in Timestamp format.
-    assert type(time_data) == pd.Timestamp
-    assert time_data == pd.Timestamp("2035-01-22 00:00:00")
+    # Checks that 'Time' data is in correct dtype
+    assert np.issubdtype(df["Time"].dtype, np.datetime64)
 
-    # Checks that all other data values are 0.
-    float_data = df.drop(columns=["Time"]).loc[0, :].values.flatten().tolist()
-    assert all(v == 0 for v in float_data)
+    # Checks that all other data types are int or float
+    dtypes = df.dtypes
+    assert ((dtypes.drop("Time") == int) | (dtypes.drop("Time") == float)).all()
 
 
 def test_append_opal_data(opal_data):
     """Tests appending new row of Opal data to Dataframe using custom accessor."""
-    from datahub.opal import OPAL_START_DATE, create_opal_frame
+    from datahub.opal import OPAL_START_DATE, create_opal_frame, opal_headers
 
     data_1 = opal_data.copy()
 
@@ -40,11 +37,10 @@ def test_append_opal_data(opal_data):
 
     # Checks that Dataframe has an additonal row each time .append is used.
     df.opal.append(data_1)
-    assert len(df.columns) == 41
-    assert len(df.index) == 2
+    assert df.shape == (1, len(opal_headers))
 
     df.opal.append(data_2)
-    assert len(df.index) == 3
+    assert len(df.index) == 2
 
     # Checks that data appended to Dataframe matches data input.
     data_1["time"] = pd.Timestamp(OPAL_START_DATE) + pd.to_timedelta(
@@ -54,18 +50,18 @@ def test_append_opal_data(opal_data):
         data_2["time"], unit="S"
     )
 
-    assert (df.iloc[1, :] == list(data_1.values())[1:]).all()
-    assert (df.iloc[2, :] == list(data_2.values())[1:]).all()
+    assert (df.loc[1] == list(data_1.values())[1:]).all()
+    assert (df.loc[2] == list(data_2.values())[1:]).all()
 
     # Checks that data overwrites existing rows if they have the same index value.
     df.opal.append(data_3)
-    assert len(df.index) == 3
+    assert len(df.index) == 2
 
     data_3["time"] = pd.Timestamp(OPAL_START_DATE) + pd.to_timedelta(
         data_3["time"], unit="S"
     )
 
-    assert (df.iloc[2, :] == list(data_3.values())[1:]).all()
+    assert (df.loc[2] == list(data_3.values())[1:]).all()
 
 
 def test_append_validate(opal_data):
@@ -88,12 +84,12 @@ def test_append_validate(opal_data):
         df.opal.append(opal_data)
 
     # Checks that Dataframe does not have an additional row
-    assert len(df.index) == 1
+    assert len(df.index) == 0
 
 
 def test_append_opal_data_array(opal_data_array):
     """Tests appending new row of Opal data using array format."""
-    from datahub.opal import OPAL_START_DATE, create_opal_frame
+    from datahub.opal import OPAL_START_DATE, create_opal_frame, opal_headers
 
     data_1 = opal_data_array.copy()
 
@@ -102,11 +98,8 @@ def test_append_opal_data_array(opal_data_array):
     # Checks that Dataframe is appended to correctly with array format data.
     df.opal.append(data_1)
 
-    assert len(df.columns) == 41
-    assert len(df.index) == 2
+    assert df.shape == (1, len(opal_headers))
 
-    del data_1[5:8]
-    del data_1[0]
-    data_1[0] = pd.Timestamp(OPAL_START_DATE) + pd.to_timedelta(data_1[0], unit="S")
+    data_1[1] = pd.Timestamp(OPAL_START_DATE) + pd.to_timedelta(data_1[1], unit="S")
 
-    assert (df.iloc[1, :] == data_1[0:]).all()
+    assert (df.loc[1] == data_1[1:]).all()

tests/test_opal_api.py

@@ -4,50 +4,55 @@
 import pytest
 
 from datahub import data as dt
-from datahub.opal import create_opal_frame
 
 
 @pytest.fixture(autouse=True)
 def reset_opal_data():
     """Pytest Fixture for resetting Opal data global variable."""
+    from datahub.opal import create_opal_frame
+
     dt.opal_df = create_opal_frame()
 
 
 def test_post_opal_api(client, opal_data):
     """Tests POSTing Opal data to API."""
+    from datahub.opal import opal_headers
+
     post_data = json.dumps(opal_data.copy())
 
     # Checks that the Opal global variable can be accessed.
-    assert len(dt.opal_df.columns) == 41
-    assert len(dt.opal_df.index) == 1
+    assert dt.opal_df.shape == (0, len(opal_headers))
 
     # Checks that a POST request can be successfully made.
     response = client.post("/opal", data=post_data)
     assert response.status_code == 200
     assert response.json() == {"message": "Data submitted successfully."}
 
     # Checks that the Opal global variable has been updated.
-    assert len(dt.opal_df.index) == 2
+    assert len(dt.opal_df.index) == 1
 
 
 def test_post_opal_api_array(client, opal_data_array):
     """Tests POSTing Opal data to API in array format."""
+    opal_data_array[5:5] = [1, 2, 3]
     post_data = json.dumps({"array": opal_data_array.copy()})
 
     # Checks that the Opal global variable has been reset.
-    assert len(dt.opal_df.index) == 1
+    assert len(dt.opal_df.index) == 0
 
     # Checks that a POST request can be successfully made.
     response = client.post("/opal", data=post_data)
     assert response.status_code == 200
     assert response.json() == {"message": "Data submitted successfully."}
 
     # Checks that the Opal global variable has been updated.
-    assert len(dt.opal_df.index) == 2
+    assert len(dt.opal_df.index) == 1
 
 
 def test_post_opal_api_invalid(client, opal_data, opal_data_array):
     """Tests error handling for invalid dict POST data."""
+    from datahub.opal import opal_headers
+
     invalid_data = opal_data.copy()
     invalid_data.pop("frame")
     post_data = json.dumps(invalid_data)
@@ -65,7 +70,7 @@ def test_post_opal_api_invalid(client, opal_data, opal_data_array):
     response = client.post("/opal", data=post_data)
     assert response.status_code == 400
     assert response.json() == {
-        "detail": "Array has invalid length. Expecting 45 items."
+        "detail": f"Array has invalid length. Expecting {len(opal_headers) + 4} items."
     }
 
     # Check that error is raised when the data on the server has been corrupted
@@ -79,7 +84,7 @@ def test_post_opal_api_invalid(client, opal_data, opal_data_array):
     }
 
     # Check that the Opal global variable has not been updated.
-    assert len(dt.opal_df.index) == 1
+    assert len(dt.opal_df.index) == 0
 
 
 def test_get_opal_api(client, opal_data):
@@ -108,7 +113,7 @@ def test_opal_api_get_query(client, opal_data):
     assert response.json()["data"]["index"] == [2, 3, 4]
 
     response = client.get("/opal?end=2")
-    assert response.json()["data"]["index"] == [0, 1, 2]
+    assert response.json()["data"]["index"] == [1, 2]
 
     response = client.get("/opal?start=1&end=3")
     assert response.json()["data"]["index"] == [1, 2, 3]