Migrate collection operations from APIHarnessV2 to CustomStorage

functions/log-event/main.py

@@ -5,11 +5,19 @@
 import uuid
 
 from crowdstrike.foundry.function import Function, Request, Response, APIError
-from falconpy import APIHarnessV2
+from falconpy import CustomStorage
 
 FUNC = Function.instance()
 
 
+def _app_headers() -> dict:
+    """Build app headers for CustomStorage construction."""
+    app_id = os.environ.get("APP_ID")
+    if app_id:
+        return {"X-CS-APP-ID": app_id}
+    return {}
+
+
 @FUNC.handler(method="POST", path="/log-event")
 def on_post(request: Request) -> Response:
     """
@@ -40,22 +48,12 @@ def on_post(request: Request) -> Response:
             "timestamp": int(time.time())
         }
 
-        # Allow setting APP_ID as an env variable for local testing
-        headers = {}
-        if os.environ.get("APP_ID"):
-            headers = {
-                "X-CS-APP-ID": os.environ.get("APP_ID")
-            }
-
-        api_client = APIHarnessV2()
+        api_client = CustomStorage(ext_headers=_app_headers())
         collection_name = "event_logs"
 
-        response = api_client.command("PutObject",
-                                      body=json_data,
-                                      collection_name=collection_name,
-                                      object_key=event_id,
-                                      headers=headers
-                                      )
+        response = api_client.PutObject(body=json_data,
+                                        collection_name=collection_name,
+                                        object_key=event_id)
 
         if response["status_code"] != 200:
             error_message = response.get("error", {}).get("message", "Unknown error")
@@ -68,17 +66,14 @@ def on_post(request: Request) -> Response:
             )
 
         # Query the collection to retrieve the event by id
-        query_response = api_client.command("SearchObjects",
-                                            filter=f"event_id:'{event_id}'",
-                                            collection_name=collection_name,
-                                            limit=5,
-                                            headers=headers
-                                            )
+        query_response = api_client.SearchObjects(filter=f"event_id:'{event_id}'",
+                                                  collection_name=collection_name,
+                                                  limit=5)
 
         return Response(
             body={
                 "stored": True,
-                "metadata": query_response.get("body").get("resources", [])
+                "metadata": query_response.get("body", {}).get("resources", [])
             },
             code=200
         )

functions/log-event/test_main.py

@@ -29,37 +29,37 @@ def setUp(self):
 
         importlib.reload(main)
 
-    @patch('main.APIHarnessV2')
+    @patch('main.CustomStorage')
     @patch('main.uuid.uuid4')
     @patch('main.time.time')
-    def test_on_post_success(self, mock_time, mock_uuid, mock_api_harness_class):
+    def test_on_post_success(self, mock_time, mock_uuid, mock_custom_storage_class):
         """Test successful POST request with valid event_data in body."""
         # Mock dependencies
         mock_uuid.return_value = MagicMock()
         mock_uuid.return_value.__str__ = MagicMock(return_value="test-event-id-123")
         mock_time.return_value = 1690123456
 
-        # Mock APIHarnessV2 instance
+        # Mock CustomStorage instance
         mock_api_instance = MagicMock()
-        mock_api_harness_class.return_value = mock_api_instance
+        mock_custom_storage_class.return_value = mock_api_instance
 
         # Mock successful PutObject response
-        mock_api_instance.command.side_effect = [
-            {  # PutObject response
-                "status_code": 200,
-                "body": {"success": True}
-            },
-            {  # SearchObjects response
-                "status_code": 200,
-                "body": {
-                    "resources": [{
-                        "event_id": "test-event-id-123",
-                        "data": {"test": "data"},
-                        "timestamp": 1690123456
-                    }]
-                }
+        mock_api_instance.PutObject.return_value = {
+            "status_code": 200,
+            "body": {"success": True}
+        }
+
+        # Mock successful SearchObjects response
+        mock_api_instance.SearchObjects.return_value = {
+            "status_code": 200,
+            "body": {
+                "resources": [{
+                    "event_id": "test-event-id-123",
+                    "data": {"test": "data"},
+                    "timestamp": 1690123456
+                }]
             }
-        ]
+        }
 
         request = Request()
         request.body = {
@@ -73,21 +73,18 @@ def test_on_post_success(self, mock_time, mock_uuid, mock_api_harness_class):
         self.assertIn("metadata", response.body)
         self.assertEqual(len(response.body["metadata"]), 1)
 
-        # Verify API calls
-        self.assertEqual(mock_api_instance.command.call_count, 2)
-
         # Verify PutObject call
-        put_call = mock_api_instance.command.call_args_list[0]
-        self.assertEqual(put_call[0][0], "PutObject")
+        mock_api_instance.PutObject.assert_called_once()
+        put_call = mock_api_instance.PutObject.call_args
         self.assertEqual(put_call[1]["collection_name"], "event_logs")
         self.assertEqual(put_call[1]["object_key"], "test-event-id-123")
         self.assertEqual(put_call[1]["body"]["event_id"], "test-event-id-123")
         self.assertEqual(put_call[1]["body"]["data"], {"test": "data", "message": "test event"})
         self.assertEqual(put_call[1]["body"]["timestamp"], 1690123456)
 
         # Verify SearchObjects call
-        search_call = mock_api_instance.command.call_args_list[1]
-        self.assertEqual(search_call[0][0], "SearchObjects")
+        mock_api_instance.SearchObjects.assert_called_once()
+        search_call = mock_api_instance.SearchObjects.call_args
         self.assertEqual(search_call[1]["filter"], "event_id:'test-event-id-123'")
         self.assertEqual(search_call[1]["collection_name"], "event_logs")
 
@@ -101,20 +98,20 @@ def test_on_post_missing_event_data(self):
         self.assertEqual(len(response.errors), 1)
         self.assertEqual(response.errors[0].message, "missing event_data")
 
-    @patch('main.APIHarnessV2')
+    @patch('main.CustomStorage')
     @patch('main.uuid.uuid4')
     @patch('main.time.time')
-    def test_on_post_put_object_error(self, mock_time, mock_uuid, mock_api_harness_class):
+    def test_on_post_put_object_error(self, mock_time, mock_uuid, mock_custom_storage_class):
         """Test POST request when PutObject API returns an error."""
         # Mock dependencies
         mock_uuid.return_value = MagicMock()
         mock_uuid.return_value.__str__ = MagicMock(return_value="test-event-id-123")
         mock_time.return_value = 1690123456
 
-        # Mock APIHarnessV2 instance with error response
+        # Mock CustomStorage instance with error response
         mock_api_instance = MagicMock()
-        mock_api_harness_class.return_value = mock_api_instance
-        mock_api_instance.command.return_value = {
+        mock_custom_storage_class.return_value = mock_api_instance
+        mock_api_instance.PutObject.return_value = {
             "status_code": 500,
             "error": {"message": "Internal server error"}
         }
@@ -130,18 +127,18 @@ def test_on_post_put_object_error(self, mock_time, mock_uuid, mock_api_harness_c
         self.assertEqual(len(response.errors), 1)
         self.assertIn("Failed to store event: Internal server error", response.errors[0].message)
 
-    @patch('main.APIHarnessV2')
+    @patch('main.CustomStorage')
     @patch('main.uuid.uuid4')
     @patch('main.time.time')
-    def test_on_post_exception_handling(self, mock_time, mock_uuid, mock_api_harness_class):
+    def test_on_post_exception_handling(self, mock_time, mock_uuid, mock_custom_storage_class):
         """Test POST request when an exception is raised."""
         # Mock dependencies
         mock_uuid.return_value = MagicMock()
         mock_uuid.return_value.__str__ = MagicMock(return_value="test-event-id-123")
         mock_time.return_value = 1690123456
 
-        # Mock APIHarnessV2 to raise an exception
-        mock_api_harness_class.side_effect = ConnectionError("Connection failed")
+        # Mock CustomStorage to raise an exception
+        mock_custom_storage_class.side_effect = ConnectionError("Connection failed")
 
         request = Request()
         request.body = {
@@ -155,23 +152,27 @@ def test_on_post_exception_handling(self, mock_time, mock_uuid, mock_api_harness
         self.assertIn("Error saving collection: Connection failed", response.errors[0].message)
 
     @patch.dict('main.os.environ', {'APP_ID': 'test-app-123'})
-    @patch('main.APIHarnessV2')
+    @patch('main.CustomStorage')
     @patch('main.uuid.uuid4')
     @patch('main.time.time')
-    def test_on_post_with_app_id_header(self, mock_time, mock_uuid, mock_api_harness_class):
+    def test_on_post_with_app_id_header(self, mock_time, mock_uuid, mock_custom_storage_class):
         """Test POST request with APP_ID environment variable set."""
         # Mock dependencies
         mock_uuid.return_value = MagicMock()
         mock_uuid.return_value.__str__ = MagicMock(return_value="test-event-id-123")
         mock_time.return_value = 1690123456
 
-        # Mock APIHarnessV2 instance
+        # Mock CustomStorage instance
         mock_api_instance = MagicMock()
-        mock_api_harness_class.return_value = mock_api_instance
-        mock_api_instance.command.side_effect = [
-            {"status_code": 200, "body": {"success": True}},
-            {"status_code": 200, "body": {"resources": []}}
-        ]
+        mock_custom_storage_class.return_value = mock_api_instance
+        mock_api_instance.PutObject.return_value = {
+            "status_code": 200,
+            "body": {"success": True}
+        }
+        mock_api_instance.SearchObjects.return_value = {
+            "status_code": 200,
+            "body": {"resources": []}
+        }
 
         request = Request()
         request.body = {
@@ -182,9 +183,10 @@ def test_on_post_with_app_id_header(self, mock_time, mock_uuid, mock_api_harness
 
         self.assertEqual(response.code, 200)
 
-        # Verify that headers with APP_ID were passed to both API calls
-        for call in mock_api_instance.command.call_args_list:
-            self.assertEqual(call[1]["headers"], {"X-CS-APP-ID": "test-app-123"})
+        # Verify that CustomStorage was constructed with ext_headers containing APP_ID
+        mock_custom_storage_class.assert_called_once_with(
+            ext_headers={"X-CS-APP-ID": "test-app-123"}
+        )
 
 
 if __name__ == "__main__":