[feat] Added count property to WorkflowSearchResponse

Author: Aryamanz29

pyatlan/model/workflow.py

@@ -234,6 +234,10 @@ def __init__(self, **data: Any):
         self._size = data.get("size")  # type: ignore[assignment]
         self._start = data.get("start")  # type: ignore[assignment]
 
+    @property
+    def count(self):
+        return self.hits.total.get("value", 0) if self.hits and self.hits.total else 0
+
     def current_page(self) -> Optional[List[WorkflowSearchResult]]:
         return self.hits.hits  # type: ignore
 

tests/integration/test_workflow_client.py

@@ -134,10 +134,10 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
     runs = client.workflow.get_runs(
         workflow_name=workflow.metadata.name, workflow_phase=AtlanWorkflowPhase.RUNNING
     )
-    assert runs
-    total_runs = runs.hits.hits  # type: ignore
-    assert len(total_runs) == 1  # type: ignore
-    run = total_runs[0]  # type: ignore
+    assert runs and runs.count == 1
+    current_page = runs.current_page()
+    assert current_page is not None and len(current_page) == 1
+    run = current_page[0]
     assert run and run.id
     assert workflow.metadata.name and (workflow.metadata.name in run.id)
 
@@ -178,7 +178,7 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
         [AtlanWorkflowPhase.FAILED], started_at="now-1h"
     )
     assert runs_status
-    workflow_run_status = runs_status.hits.hits[0]  # type: ignore
+    workflow_run_status = runs_status.current_page()[0]  # type: ignore
     start_time = workflow_run_status.source.status.startedAt  # type: ignore
     start_datetime = datetime.strptime(start_time, "%Y-%m-%dT%H:%M:%SZ")  # type: ignore
     start_datetime = start_datetime.replace(tzinfo=timezone.utc)

tests/unit/test_client.py

@@ -2563,9 +2563,9 @@ def test_get_all_pagation(group_client, mock_api_caller):
     ]
 
     groups = group_client.get_all(limit=2)
-    assert len(groups.records) == 2
-    assert groups.records[0].id == "1"
-    assert groups.records[1].id == "2"
+    assert len(groups.current_page()) == 2
+    assert groups.current_page()[0].id == "1"
+    assert groups.current_page()[1].id == "2"
     assert mock_api_caller._call_api.call_count == 1
     mock_api_caller.reset_mock()
 
@@ -2577,7 +2577,7 @@ def test_get_all_empty_response_with_raw_records(group_client, mock_api_caller):
     ]
 
     groups = group_client.get_all()
-    assert len(groups.records) == 0
+    assert len(groups.current_page()) == 0
     mock_api_caller.reset_mock()
 
 
@@ -2593,9 +2593,9 @@ def test_get_all_with_columns(group_client, mock_api_caller):
     columns = ["alias"]
     groups = group_client.get_all(limit=10, columns=columns)
 
-    assert len(groups.records) == 2
-    assert groups.records[0].id == "1"
-    assert groups.records[0].alias == "Group1"
+    assert len(groups.current_page()) == 2
+    assert groups.current_page()[0].id == "1"
+    assert groups.current_page()[0].alias == "Group1"
     mock_api_caller._call_api.assert_called_once()
     query_params = mock_api_caller._call_api.call_args.kwargs["query_params"]
     assert query_params["columns"] == columns
@@ -2613,9 +2613,9 @@ def test_get_all_sorting(group_client, mock_api_caller):
 
     groups = group_client.get_all(limit=10, sort="alias")
 
-    assert len(groups.records) == 2
-    assert groups.records[0].id == "1"
-    assert groups.records[0].alias == "Group1"
+    assert len(groups.current_page()) == 2
+    assert groups.current_page()[0].id == "1"
+    assert groups.current_page()[0].alias == "Group1"
     mock_api_caller._call_api.assert_called_once()
     query_params = mock_api_caller._call_api.call_args.kwargs["query_params"]
     assert query_params["sort"] == "alias"

tests/unit/test_workflow_client.py

@@ -273,23 +273,19 @@ def test_find_by_type(client: WorkflowClient, mock_api_caller):
 
 
 def test_find_runs_by_status_and_time_range(client: WorkflowClient, mock_api_caller):
-    raw_json = {"shards": {"dummy": None}, "hits": {"total": {"dummy": None}}}
+    raw_json = {"_shards": {"dummy": None}, "hits": {"total": {"dummy": None}}}
     mock_api_caller._call_api.return_value = raw_json
 
     status = [AtlanWorkflowPhase.SUCCESS, AtlanWorkflowPhase.FAILED]
     started_at = "now-2h"
     finished_at = "now-1h"
-
-    assert (
-        client.find_runs_by_status_and_time_range(
-            status=status,
-            started_at=started_at,
-            finished_at=finished_at,
-            from_=10,
-            size=5,
-        )
-        == []
-    )
+    assert client.find_runs_by_status_and_time_range(
+        status=status,
+        started_at=started_at,
+        finished_at=finished_at,
+        from_=10,
+        size=5,
+    ) == WorkflowSearchResponse(**raw_json)
     mock_api_caller._call_api.assert_called_once()
     assert isinstance(
         mock_api_caller._call_api.call_args.kwargs["request_obj"], WorkflowSearchRequest
@@ -556,14 +552,13 @@ def test_workflow_get_runs(
     mock_api_caller,
     search_response: WorkflowSearchResponse,
 ):
-    mock_api_caller._call_api.return_value = search_response.dict()
+    mock_api_caller._call_api.return_value = search_response.dict(by_alias=True)
     response = client.get_runs(
         workflow_name="test-workflow",
         workflow_phase=AtlanWorkflowPhase.RUNNING,
     )
 
-    assert search_response.hits
-    assert response == search_response.hits.hits
+    assert response == WorkflowSearchResponse(**search_response.dict())
     assert mock_api_caller._call_api.call_count == 1
     mock_api_caller.reset_mock()