added produce test with duplicates

alercebroker · Dec 2, 2020 · 455dc19 · 455dc19
1 parent 0f72750
commit 455dc19
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/features/step.py b/features/step.py
@@ -180,7 +180,9 @@ def update_db(self, to_update, out_columns, apply_get_fid):
         to_update.columns = out_columns
         to_update["fid"] = to_update["name"].apply(apply_get_fid)
         to_update["version"] = self.feature_version.version
-        to_update["name"] = to_update["name"].apply(lambda x: self.check_feature_name(x))
+        to_update["name"] = to_update["name"].apply(
+            lambda x: self.check_feature_name(x)
+        )
         to_update.rename(
             columns={
                 "oid": "_oid",

diff --git a/tests/unittest/test_step.py b/tests/unittest/test_step.py
@@ -268,7 +268,9 @@ def test_check_feature_name(self):
     @mock.patch.object(FeaturesComputer, "produce")
     @mock.patch.object(FeaturesComputer, "add_to_db")
     @mock.patch.object(FeaturesComputer, "insert_feature_version")
-    def test_execute_with_producer(self, insert_feature_version, add_to_db, produce, compute_features):
+    def test_execute_with_producer(
+        self, insert_feature_version, add_to_db, produce, compute_features
+    ):
         messages = [
             {
                 "oid": "ZTF1",
@@ -290,7 +292,9 @@ def test_execute_with_producer(self, insert_feature_version, add_to_db, produce,
     @mock.patch.object(FeaturesComputer, "produce")
     @mock.patch.object(FeaturesComputer, "add_to_db")
     @mock.patch.object(FeaturesComputer, "insert_feature_version")
-    def test_execute_duplicates(self, insert_feature_version, add_to_db, produce, compute_features):
+    def test_execute_duplicates(
+        self, insert_feature_version, add_to_db, produce, compute_features
+    ):
         message1 = {
             "oid": "ZTF1",
             "candid": 123,
@@ -314,3 +318,15 @@ def test_execute_duplicates(self, insert_feature_version, add_to_db, produce, co
         mock_feature_version.version = self.step_config["FEATURE_VERSION"]
         self.step.execute(messages)
         produce.assert_called_once()
+
+    def test_produce(self):
+        alert_data = pd.DataFrame({"oid": ["OID", "OID"], "candid": [123, 123]})
+        features = pd.DataFrame({"oid": ["OID"], "feature1": 1, "feature2": 2})
+        features.set_index("oid", inplace=True)
+        self.step.produce(features, alert_data)
+        expected_message = {
+            "features": {"feature1": 1, "feature2": 2},
+            "oid": "OID",
+            "candid": 123,
+        }
+        self.step.producer.produce.assert_called_with(expected_message, key="OID")