From 2bd22a92c9f18a53212f71fefefe6fddf23311fe Mon Sep 17 00:00:00 2001
From: bubriks <bubriks@gmail.com>
Date: Fri, 13 Dec 2024 11:04:22 +0200
Subject: [PATCH] add unit test for get_headers

---
 python/tests/core/test_kafka_engine.py        | 38 ++++++++++++++++++-
 python/tests/fixtures/backend_fixtures.py     |  1 +
 .../fixtures/online_ingestion_fixtures.json   | 22 +++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 python/tests/fixtures/online_ingestion_fixtures.json

diff --git a/python/tests/core/test_kafka_engine.py b/python/tests/core/test_kafka_engine.py
index 88085689e..51344e2cb 100644
--- a/python/tests/core/test_kafka_engine.py
+++ b/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'
+        }
diff --git a/python/tests/fixtures/backend_fixtures.py b/python/tests/fixtures/backend_fixtures.py
index dd455b699..71a55548f 100644
--- a/python/tests/fixtures/backend_fixtures.py
+++ b/python/tests/fixtures/backend_fixtures.py
@@ -71,6 +71,7 @@
     "transformer",
     "user",
     "validation_report",
+    "online_ingestion",
 ]
 
 backend_fixtures_json = {}
diff --git a/python/tests/fixtures/online_ingestion_fixtures.json b/python/tests/fixtures/online_ingestion_fixtures.json
new file mode 100644
index 000000000..23d804e60
--- /dev/null
+++ b/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": []
+    }
+  }
+}