Migrate collection operations from APIHarnessV2 to CustomStorage (#118)

mraible · web-flow · commit 733d2093b707 · 2026-04-15T22:08:17.000-06:00
diff --git a/e2e/src/pages/WorkflowsPage.ts b/e2e/src/pages/WorkflowsPage.ts
@@ -246,4 +246,81 @@ export class WorkflowsPage extends BasePage {
       `Execute and verify workflow: ${workflowName}`
     );
   }
+
+  /**
+   * Check the actual execution status by viewing the execution details.
+   * Navigates to the execution log, waits for the execution to complete,
+   * expands the execution row, and checks the status.
+   */
+  async verifyWorkflowExecutionCompleted(timeoutMs = 120000): Promise<void> {
+    return this.withTiming(
+      async () => {
+        this.logger.info('Checking workflow execution status in detail view');
+
+        // The "View" link opens in a new tab - capture it
+        const viewLink = this.page.getByRole('link', { name: /^view$/i });
+        await viewLink.waitFor({ state: 'visible', timeout: 10000 });
+
+        const [executionPage] = await Promise.all([
+          this.page.context().waitForEvent('page'),
+          viewLink.click(),
+        ]);
+
+        // Wait for the new tab to load (execution pages can be slow to render)
+        await executionPage.waitForLoadState('networkidle');
+        await executionPage.waitForLoadState('domcontentloaded');
+        this.logger.info('Execution page opened in new tab');
+
+        // Wait for "Execution status" to appear (proves execution details loaded)
+        const statusLabel = executionPage.getByText('Execution status');
+        await statusLabel.waitFor({ state: 'visible', timeout: 60000 });
+        this.logger.info('Execution details visible');
+
+        // Poll until execution reaches a terminal state
+        this.logger.info(`Waiting up to ${timeoutMs / 1000}s for execution to complete...`);
+
+        const startTime = Date.now();
+        while (Date.now() - startTime < timeoutMs) {
+          // Re-find status label each iteration (DOM recreated on reload)
+          const currentStatusLabel = executionPage.getByText('Execution status');
+          await currentStatusLabel.waitFor({ state: 'visible', timeout: 15000 });
+          const statusContainer = currentStatusLabel.locator('..');
+          const statusText = await statusContainer.textContent() || '';
+          const currentStatus = statusText.replace('Execution status', '').trim();
+          this.logger.info(`Current status: ${currentStatus}`);
+
+          if (currentStatus.toLowerCase().includes('failed')) {
+            // Capture error details
+            const pageContent = await executionPage.textContent('body') || '';
+            const messageMatch = pageContent.match(/"message":\s*"([^"]+)"/);
+
+            let errorMessage = 'Workflow action failed';
+            if (messageMatch) {
+              errorMessage = messageMatch[1];
+            }
+
+            await executionPage.close();
+            this.logger.error(`Workflow execution failed: ${errorMessage}`);
+            throw new Error(`Workflow execution failed: ${errorMessage}`);
+          }
+
+          if (!currentStatus.toLowerCase().includes('in progress')) {
+            // Terminal state that isn't "Failed"
+            await executionPage.close();
+            this.logger.success(`Workflow execution completed with status: ${currentStatus}`);
+            return;
+          }
+
+          await executionPage.waitForTimeout(5000);
+
+          // Reload to get updated status - the page doesn't auto-refresh
+          await executionPage.reload({ waitUntil: 'networkidle' });
+        }
+
+        await executionPage.close();
+        throw new Error('Workflow execution timed out - still in progress');
+      },
+      'Verify workflow execution completed'
+    );
+  }
 }
diff --git a/e2e/tests/foundry.spec.ts b/e2e/tests/foundry.spec.ts
@@ -1,4 +1,4 @@
-import { test, expect } from '../src/fixtures';
+import { test } from '../src/fixtures';
 
 test.describe.configure({ mode: 'serial' });
 
@@ -9,16 +9,21 @@ test.describe('Functions with Python - E2E Tests', () => {
   });
 
   test('should execute Test hello function workflow', async ({ workflowsPage }) => {
+    test.setTimeout(180000);
     await workflowsPage.navigateToWorkflows();
     await workflowsPage.executeAndVerifyWorkflow('Test hello function');
+    await workflowsPage.verifyWorkflowExecutionCompleted();
   });
 
   test('should execute Test log-event function workflow', async ({ workflowsPage }) => {
+    test.setTimeout(180000);
     await workflowsPage.navigateToWorkflows();
     await workflowsPage.executeAndVerifyWorkflow('Test log-event function');
+    await workflowsPage.verifyWorkflowExecutionCompleted();
   });
 
   test('should execute Test host-details function workflow', async ({ workflowsPage, hostManagementPage }) => {
+    test.setTimeout(180000);
     // Get first available host ID
     const hostId = await hostManagementPage.getFirstHostId();
 
@@ -32,6 +37,7 @@ test.describe('Functions with Python - E2E Tests', () => {
     await workflowsPage.executeAndVerifyWorkflow('Test host-details function', {
       'Host ID': hostId
     });
+    await workflowsPage.verifyWorkflowExecutionCompleted();
   });
 
   test('should render Test servicenow function workflow (without execution)', async ({ workflowsPage }) => {
diff --git a/functions/log-event/main.py b/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()
+        custom_storage = 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 = custom_storage.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 = custom_storage.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
         )
diff --git a/functions/log-event/test_main.py b/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__":
