MRG: Merge pull request #14 from octue/update-attribute-handling

Update attribute handling

Author: cortadocodes

functions/event_handler/main.py

@@ -62,7 +62,7 @@ def handle_event(cloud_event):
         "sender_sdk_version": attributes.pop("sender_sdk_version"),
         "recipient": attributes.pop("recipient"),
         "question_uuid": attributes.pop("question_uuid"),
-        "parent_question_uuid": attributes.pop("parent_question_uuid"),
+        "parent_question_uuid": attributes.pop("parent_question_uuid", None),
         "originator_question_uuid": attributes.pop("originator_question_uuid"),
         # Backend-specific metadata.
         "backend": BACKEND,

poetry.lock

@@ -1,6 +1,6 @@
 [project]
 name = "twined-gcp"
-version = "0.7.0"
+version = "0.7.1"
 description = ""
 authors = [
     {name = "Marcus Lugg", email = "[email protected]"}

pyproject.toml

@@ -24,7 +24,12 @@
     "question_uuid": QUESTION_UUID,
     "parent_question_uuid": "1d897229-155d-498d-b6ae-21960fab3754",
     "originator_question_uuid": "fb6cf9a3-84fb-45ce-a4da-0d2257bec319",
-    "forward_logs": True,
+    "retry_count": "0",
+    "forward_logs": "1",
+    "save_diagnostics": "SAVE_DIAGNOSTICS_ON_CRASH",
+    "cpus": "1",
+    "memory": "2Gi",
+    "ephemeral_storage": "256Mi",
 }
 
 
@@ -75,9 +80,16 @@ def test_store_pub_sub_event_in_bigquery(self):
                 "datetime": "2024-04-11T09:26:39.144818",
                 "uuid": "c8bda9fa-f072-4330-92b1-96920d06b28d",
                 "kind": "heart",
-                "event": {"some": "data"},
+                "event": {
+                    "some": "data",
+                },
                 "other_attributes": {
-                    "forward_logs": True,
+                    "retry_count": "0",
+                    "forward_logs": "1",
+                    "save_diagnostics": "SAVE_DIAGNOSTICS_ON_CRASH",
+                    "cpus": "1",
+                    "memory": "2Gi",
+                    "ephemeral_storage": "256Mi",
                 },
                 "parent": "octue/parent-test-service:5.6.3",
                 "originator": "octue/ancestor-test-service:5.6.3",
@@ -93,13 +105,51 @@ def test_store_pub_sub_event_in_bigquery(self):
             },
         )
 
+    def test_store_pub_sub_event_in_bigquery_with_no_parent_question_uuid(self):
+        """Test that an event with no `parent_question_uuid` attribute is handled correctly."""
+        attributes = copy.copy(EVENT_ATTRIBUTES)
+        attributes.pop("parent_question_uuid")
+
+        cloud_event = MockCloudEvent(
+            data={
+                "message": {
+                    "data": base64.b64encode(b'{"kind": "heart", "some": "data"}'),
+                    "attributes": attributes,
+                    "messageId": "1234",
+                }
+            }
+        )
+
+        mock_big_query_client = MockBigQueryClient()
+
+        with patch("functions.event_handler.main.BigQueryClient", return_value=mock_big_query_client):
+            handle_event(cloud_event)
+
+        self.assertIsNone(mock_big_query_client.inserted_rows[0][0]["parent_question_uuid"])
+
     def test_question_is_dispatched_to_kueue(self):
         """Test that questions are dispatched to Kueue correctly."""
+        event_attributes = {
+            "datetime": "2024-04-11T09:26:39.144818",
+            "uuid": "c8bda9fa-f072-4330-92b1-96920d06b28d",
+            "parent": "octue/parent-test-service:5.6.3",
+            "originator": "octue/ancestor-test-service:5.6.3",
+            "sender": "octue/test-service:5.6.3",
+            "sender_type": "PARENT",
+            "sender_sdk_version": "1.0.3",
+            "recipient": SRUID,
+            "question_uuid": QUESTION_UUID,
+            "parent_question_uuid": "1d897229-155d-498d-b6ae-21960fab3754",
+            "originator_question_uuid": "fb6cf9a3-84fb-45ce-a4da-0d2257bec319",
+            "retry_count": "0",
+            "forward_logs": "1",
+        }
+
         cloud_event = MockCloudEvent(
             data={
                 "message": {
                     "data": base64.b64encode(b'{"kind": "question", "input_values": {"some": "data"}}'),
-                    "attributes": copy.copy(EVENT_ATTRIBUTES),
+                    "attributes": copy.copy(event_attributes),
                     "messageId": "1234",
                 }
             }
@@ -118,11 +168,13 @@ def test_question_is_dispatched_to_kueue(self):
         self.assertEqual(container.name, job.metadata.name)
         self.assertEqual(container.image, f"some-artifact-registry-url/{SRUID}")
         self.assertEqual(container.command, ["octue", "question", "ask", "local"])
+
+        # Check the default resource requirements are used.
         self.assertEqual(container.resources, {"requests": {"cpu": 1, "ephemeral-storage": "1Gi", "memory": "500Mi"}})
 
         self.assertEqual(
             container.args,
-            ["--attributes", json.dumps(EVENT_ATTRIBUTES), "--input-values", '{"some": "data"}'],
+            ["--attributes", json.dumps(event_attributes), "--input-values", '{"some": "data"}'],
         )
 
         environment_variables = [variable.to_dict() for variable in container.env]