Merge get_schema functionality into assemble_schema

Author: Fireye04

src/sasquatchbackpack/sasquatch.py

@@ -46,11 +46,7 @@ def __init__(self, topic_name: str, *, uses_redis: bool) -> None:
         self.uses_redis = uses_redis
 
     @abstractmethod
-    def get_schema(self, namespace: str) -> AvroBaseModel:
-        pass
-
-    @abstractmethod
-    def assemble_schema(self, records: dict, namespace: str) -> AvroBaseModel:
+    def assemble_schema(self, namespace: str, records: dict) -> AvroBaseModel:
         pass
 
     @abstractmethod
@@ -184,7 +180,7 @@ async def _dispatch(
 
     for record in records:
         avro: bytes = await schema_manager.serialize(
-            source.assemble_schema(record["value"], namespace)
+            source.assemble_schema(namespace, record["value"])
         )
         headers = {
             "content-type": "avro",
@@ -220,7 +216,7 @@ def __init__(
         self.source = source
         self.config = DispatcherConfig()
 
-        self.schema = self.source.get_schema(self.config.namespace)
+        self.schema = self.source.assemble_schema(self.config.namespace)
         nest_asyncio.apply()
         self.redis = RedisManager(
             self.config.redis_address

src/sasquatchbackpack/sources/usgs/scripts.py

@@ -143,28 +143,29 @@ def get_records(self) -> list[dict]:
                 f"A connection error occurred while fetching records: {ce}"
             ) from ce
 
-    def assemble_schema(self, record: dict, namespace: str) -> AvroBaseModel:
-        schema = {
-            "timestamp": record["timestamp"],
-            "id": record["id"],
-            "latitude": record["latitude"],
-            "longitude": record["longitude"],
-            "depth": record["depth"],
-            "magnitude": record["magnitude"],
-            "namespace": namespace,
-        }
-        return EarthquakeSchema.parse_obj(data=schema)
-
-    def get_schema(self, namespace: str) -> AvroBaseModel:
-        schema = {
-            "timestamp": 1,
-            "id": "default",
-            "latitude": 1.0,
-            "longitude": 1.0,
-            "depth": 1.0,
-            "magnitude": 1.0,
-            "namespace": namespace,
-        }
+    def assemble_schema(
+        self, namespace: str, record: dict = {}
+    ) -> AvroBaseModel:
+        if record == {}:
+            schema = {
+                "timestamp": 1,
+                "id": "default",
+                "latitude": 1.0,
+                "longitude": 1.0,
+                "depth": 1.0,
+                "magnitude": 1.0,
+                "namespace": namespace,
+            }
+        else:
+            schema = {
+                "timestamp": record["timestamp"],
+                "id": record["id"],
+                "latitude": record["latitude"],
+                "longitude": record["longitude"],
+                "depth": record["depth"],
+                "magnitude": record["magnitude"],
+                "namespace": namespace,
+            }
         return EarthquakeSchema.parse_obj(data=schema)
 
     def get_redis_key(self, datapoint: dict) -> str:

tests/dispatcher_test.py

@@ -30,18 +30,19 @@ def __init__(self, current_records: list[dict[str, str]]) -> None:
         )
         self.records = current_records
 
-    def assemble_schema(self, record: dict, namespace: str) -> AvroBaseModel:
-        schema = {
-            "id": record["id"],
-            "namespace": namespace,
-        }
-        return TestSchema.parse_obj(data=schema)
-
-    def get_schema(self, namespace: str) -> AvroBaseModel:
-        schema = {
-            "id": "default",
-            "namespace": namespace,
-        }
+    def assemble_schema(
+        self, namespace: str, record: dict = {}
+    ) -> AvroBaseModel:
+        if record == {}:
+            schema = {
+                "id": "default",
+                "namespace": namespace,
+            }
+        else:
+            schema = {
+                "id": record["id"],
+                "namespace": namespace,
+            }
         return TestSchema.parse_obj(data=schema)
 
     def get_records(self) -> list[dict]: