add unit test for get_headers

python/tests/core/test_kafka_engine.py

@@ -15,8 +15,8 @@
 #
 import importlib
 
-from hsfs import storage_connector
-from hsfs.core import constants, kafka_engine
+from hsfs import storage_connector, feature_group
+from hsfs.core import constants, kafka_engine, online_ingestion
 
 
 if constants.HAS_CONFLUENT_KAFKA:
@@ -523,3 +523,37 @@ def test_spark_get_kafka_config_internal_kafka(self, mocker, backend_fixtures):
             mock_storage_connector_api.return_value.get_kafka_connector.call_args[0][1]
             is False
         )
+
+    def test_get_headers(self, mocker, backend_fixtures):
+        # Arrange
+        mocker.patch("hopsworks_common.client.get_instance")
+        mock_online_ingestion_api = mocker.patch(
+            "hsfs.core.online_ingestion_api.OnlineIngestionApi"
+        )
+        json = backend_fixtures["online_ingestion"]["get"]["response"]
+        oi = online_ingestion.OnlineIngestion.from_response_json(json)
+        mock_online_ingestion_api.return_value.create_online_ingestion.return_value = oi
+
+        fg = feature_group.FeatureGroup(
+            id=111,
+            name="test",
+            version=1,
+            featurestore_id=99,
+        )
+        fg.feature_store = mocker.Mock()
+        fg.feature_store.project_id = 234
+
+        fg._subject = {"id": 823}
+
+        # Act
+        results = kafka_engine.get_headers(
+            fg, num_entries=10
+        )
+
+        # Assert
+        assert results == {
+            'featureGroupId': b'111',
+            'onlineIngestionId': b'1',
+            'projectId': b'234',
+            'subjectId': b'823'
+        }

python/tests/fixtures/backend_fixtures.py

@@ -71,6 +71,7 @@
     "transformer",
     "user",
     "validation_report",
+    "online_ingestion",
 ]
 
 backend_fixtures_json = {}

python/tests/fixtures/online_ingestion_fixtures.json

@@ -0,0 +1,22 @@
+{
+  "get": {
+    "response": {
+      "count": 1,
+      "items": [
+        {
+          "id": 1,
+          "num_entries": 10,
+          "processed_entries": 8,
+          "inserted_entries": 6,
+          "aborted_entries": 2
+        }
+      ]
+    }
+  },
+  "get_empty": {
+    "response": {
+      "count": 0,
+      "items": []
+    }
+  }
+}