Added new method find_by_status_and_interval() along with it's unit and integration tests

Author: vaibhavatlan

pyatlan/client/workflow.py

@@ -24,7 +24,17 @@
 )
 from pyatlan.errors import ErrorCode
 from pyatlan.model.enums import AtlanWorkflowPhase, WorkflowPackage
-from pyatlan.model.search import Bool, NestedQuery, Prefix, Query, Regexp, Term
+from pyatlan.model.search import (
+    Bool,
+    Exists,
+    NestedQuery,
+    Prefix,
+    Query,
+    Range,
+    Regexp,
+    Term,
+    Terms,
+)
 from pyatlan.model.workflow import (
     ReRunRequest,
     ScheduleQueriesSearchRequest,
@@ -149,6 +159,43 @@ def find_run_by_id(self, id: str) -> Optional[WorkflowSearchResult]:
         response = self._find_runs(query, size=1)
         return results[0] if (results := response.hits and response.hits.hits) else None
 
+    @validate_arguments
+    def find_by_status_and_interval(
+        self, status: List[AtlanWorkflowPhase], interval: int
+    ) -> List[WorkflowSearchResult]:
+        """
+        Find workflows based on their status and interval
+
+        :status: list of the status of the workflows to filter with
+        :interval: time interval in hours to search for workflows
+        :returns: the list of workflows of the provided type, with the most-recently created first
+        :raises ValidationError: If the provided status is an invalid AtlanWorkflowPhase
+        :raises AtlanError: on any API communication issue
+
+        """
+
+        run_lookup_query = Bool(
+            must=[
+                NestedQuery(
+                    query=Terms(
+                        field="metadata.labels.workflows.argoproj.io/phase.keyword",
+                        values=[state.value for state in status],
+                    ),
+                    path="metadata",
+                ),
+                Range(field="status.finishedAt", gt=f"now-{interval}h"),
+                NestedQuery(
+                    query=Exists(field="metadata.labels.workflows.argoproj.io/creator"),
+                    path="metadata",
+                ),
+            ],
+            must_not=[NestedQuery(query=Exists(field="spec.shutdown"), path="spec")],
+        )
+
+        run_lookup_results = self._find_runs(run_lookup_query)
+
+        return run_lookup_results.hits and run_lookup_results.hits.hits or []
+
     @validate_arguments
     def _find_latest_run(self, workflow_name: str) -> Optional[WorkflowSearchResult]:
         """

tests/integration/test_workflow_client.py

@@ -162,6 +162,19 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
     workflow_status = client.workflow.monitor(workflow_name=workflow_name)
     assert workflow_status == AtlanWorkflowPhase.FAILED
 
+    # Test find run by status and interval
+    runs_status = client.workflow.find_by_status_and_interval(
+        [AtlanWorkflowPhase.FAILED], 1
+    )
+    assert runs_status
+    workflow_run_status = runs_status[0]
+    assert (
+        workflow_run_status
+        and workflow_run_status.source
+        and workflow_run_status.source.status
+        and workflow_run_status.source.status.phase == AtlanWorkflowPhase.FAILED
+    )
+
     # Test find run by id
     workflow_run = client.workflow.find_run_by_id(id=run.id)
     assert (

tests/unit/test_workflow_client.py

@@ -272,6 +272,20 @@ def test_find_by_type(client: WorkflowClient, mock_api_caller):
     )
 
 
+def test_find_by_status_and_interval(client: WorkflowClient, mock_api_caller):
+    raw_json = {"shards": {"dummy": None}, "hits": {"total": {"dummy": None}}}
+    mock_api_caller._call_api.return_value = raw_json
+
+    status = [AtlanWorkflowPhase.SUCCESS, AtlanWorkflowPhase.FAILED]
+    interval = 50
+
+    assert client.find_by_status_and_interval(status, interval) == []
+    mock_api_caller._call_api.assert_called_once()
+    assert isinstance(
+        mock_api_caller._call_api.call_args.kwargs["request_obj"], WorkflowSearchRequest
+    )
+
+
 def test_find_by_id(
     client: WorkflowClient, search_response: WorkflowSearchResponse, mock_api_caller
 ):