diff --git a/features/step.py b/features/step.py index 4e6ecc5..0cd3860 100644 --- 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 index 2740a38..ed2e7c9 100644 --- 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")