From 480e7341cfb5a060be6ae5f837a0d3a0285b7293 Mon Sep 17 00:00:00 2001
From: "ddefreitas.exe" <1959720+darrell-d@users.noreply.github.com>
Date: Tue, 10 Mar 2026 16:15:53 -0400
Subject: [PATCH]  Fix workflow client to extract packageIds from dataSources

---
 processor/clients/workflow_client.py |  6 +++++-
 tests/test_workflow_client.py        | 10 +++++-----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/processor/clients/workflow_client.py b/processor/clients/workflow_client.py
index f284c8c..ecbcaf3 100644
--- a/processor/clients/workflow_client.py
+++ b/processor/clients/workflow_client.py
@@ -34,10 +34,14 @@ def get_workflow_instance(self, workflow_instance_id):
             response.raise_for_status()
             data = response.json()
 
+            package_ids = []
+            for source in data.get("dataSources", {}).values():
+                package_ids.extend(source.get("packageIds", []))
+
             workflow_instance = WorkflowInstance(
                 id=data["uuid"],
                 dataset_id=data["datasetId"],
-                package_ids=data["packageIds"],
+                package_ids=package_ids,
             )
 
             return workflow_instance
diff --git a/tests/test_workflow_client.py b/tests/test_workflow_client.py
index fd52f2e..6be87c0 100644
--- a/tests/test_workflow_client.py
+++ b/tests/test_workflow_client.py
@@ -50,7 +50,7 @@ def test_get_workflow_instance_success(self, mock_session_manager):
         responses.add(
             responses.GET,
             "https://api.test.com/compute/workflows/runs/wf-instance-123",
-            json={"uuid": "wf-instance-123", "datasetId": "dataset-456", "packageIds": ["pkg-1", "pkg-2", "pkg-3"]},
+            json={"uuid": "wf-instance-123", "datasetId": "dataset-456", "dataSources": {"source-1": {"packageIds": ["pkg-1", "pkg-2", "pkg-3"]}}},
             status=200,
         )
 
@@ -68,7 +68,7 @@ def test_get_workflow_instance_includes_auth_header(self, mock_session_manager):
         responses.add(
             responses.GET,
             "https://api.test.com/compute/workflows/runs/wf-123",
-            json={"uuid": "wf-123", "datasetId": "ds-1", "packageIds": []},
+            json={"uuid": "wf-123", "datasetId": "ds-1", "dataSources": {}},
             status=200,
         )
 
@@ -106,7 +106,7 @@ def test_get_workflow_instance_with_single_package(self, mock_session_manager):
         responses.add(
             responses.GET,
             "https://api.test.com/compute/workflows/runs/wf-123",
-            json={"uuid": "wf-123", "datasetId": "ds-1", "packageIds": ["single-pkg"]},
+            json={"uuid": "wf-123", "datasetId": "ds-1", "dataSources": {"source-1": {"packageIds": ["single-pkg"]}}},
             status=200,
         )
 
@@ -131,7 +131,7 @@ def test_get_workflow_instance_retries_on_401(self, mock_session_manager):
         responses.add(
             responses.GET,
             "https://api.test.com/compute/workflows/runs/wf-123",
-            json={"uuid": "wf-123", "datasetId": "ds-1", "packageIds": []},
+            json={"uuid": "wf-123", "datasetId": "ds-1", "dataSources": {}},
             status=200,
         )
 
@@ -152,7 +152,7 @@ def test_get_workflow_instance_retries_on_403(self, mock_session_manager):
         responses.add(
             responses.GET,
             "https://api.test.com/compute/workflows/runs/wf-123",
-            json={"uuid": "wf-123", "datasetId": "ds-1", "packageIds": []},
+            json={"uuid": "wf-123", "datasetId": "ds-1", "dataSources": {}},
             status=200,
         )