refactor(pems_data): more generic services subpackage

the concept of a Service here is for the "what" of accessing
specific data for business logic

refactor the StationsBucket to a StationsService
StationsService needs an IDataSource to read from (the "how")

Author: thekaveman

pems_data/src/pems_data/services/__init__.py

@@ -1,13 +1,16 @@
 import pandas as pd
 
-from pems_data.s3 import S3Bucket
+from pems_data.sources.base import IDataSource
 
 
-class StationsBucket(S3Bucket):
-    """Station-specific bucket data."""
+class StationsService:
+    """Manages fetching of station-related data."""
 
-    imputation_detector_agg_5min = "imputation/detector_imputed_agg_five_minutes"
-    metadata_file = "geo/current_stations.parquet"
+    _imputation_detector_agg_5min = "imputation/detector_imputed_agg_five_minutes"
+    _metadata_file = "geo/current_stations.parquet"
+
+    def __init__(self, data_source: IDataSource):
+        self.data_source = data_source
 
     def get_district_metadata(self, district_number: str) -> pd.DataFrame:
         """Loads metadata for all stations in the selected District from S3."""
@@ -30,7 +33,7 @@ def get_district_metadata(self, district_number: str) -> pd.DataFrame:
         ]
         filters = [("DISTRICT", "=", district_number)]
 
-        return self.read_parquet(self.metadata_file, columns=columns, filters=filters)
+        return self.data_source.read(self._metadata_file, columns=columns, filters=filters)
 
     def get_imputed_agg_5min(self, station_id: str) -> pd.DataFrame:
         """Loads imputed aggregate 5 minute data for a specific station."""
@@ -45,4 +48,4 @@ def get_imputed_agg_5min(self, station_id: str) -> pd.DataFrame:
         ]
         filters = [("STATION_ID", "=", station_id)]
 
-        return self.read_parquet(self.imputation_detector_agg_5min, columns=columns, filters=filters)
+        return self.data_source.read(self._imputation_detector_agg_5min, columns=columns, filters=filters)

pems_data/src/pems_data/stations.py renamed to pems_data/src/pems_data/services/stations.py

@@ -0,0 +1,72 @@
+import pandas as pd
+import pytest
+
+from pems_data.sources.base import IDataSource
+from pems_data.services.stations import StationsService
+
+
+class TestStationsService:
+    @pytest.fixture
+    def df(self):
+        return pd.DataFrame({"STATION_ID": [1]})
+
+    @pytest.fixture
+    def data_source(self, df, mocker):
+        mock = mocker.Mock(spec=IDataSource)
+        mock.read.return_value = df
+        return mock
+
+    @pytest.fixture
+    def service(self, data_source):
+        return StationsService(data_source)
+
+    def test_imputation_detector_agg_5min(self):
+        assert StationsService._imputation_detector_agg_5min == "imputation/detector_imputed_agg_five_minutes"
+
+    def test_metadata_file(self):
+        assert StationsService._metadata_file == "geo/current_stations.parquet"
+
+    def test_get_district_metadata(self, service: StationsService, data_source: IDataSource, df):
+        district_number = "7"
+        result = service.get_district_metadata(district_number)
+
+        data_source.read.assert_called_once_with(
+            StationsService._metadata_file,
+            columns=[
+                "STATION_ID",
+                "NAME",
+                "PHYSICAL_LANES",
+                "STATE_POSTMILE",
+                "ABSOLUTE_POSTMILE",
+                "LATITUDE",
+                "LONGITUDE",
+                "LENGTH",
+                "STATION_TYPE",
+                "DISTRICT",
+                "FREEWAY",
+                "DIRECTION",
+                "COUNTY_NAME",
+                "CITY_NAME",
+            ],
+            filters=[("DISTRICT", "=", district_number)],
+        )
+        pd.testing.assert_frame_equal(result, df)
+
+    def test_get_imputed_agg_5min(self, service: StationsService, data_source: IDataSource, df):
+        station_id = "123"
+
+        result = service.get_imputed_agg_5min(station_id)
+
+        data_source.read.assert_called_once_with(
+            StationsService._imputation_detector_agg_5min,
+            columns=[
+                "STATION_ID",
+                "LANE",
+                "SAMPLE_TIMESTAMP",
+                "VOLUME_SUM",
+                "SPEED_FIVE_MINS",
+                "OCCUPANCY_AVG",
+            ],
+            filters=[("STATION_ID", "=", station_id)],
+        )
+        pd.testing.assert_frame_equal(result, df)