Major changes to the method

Author: vaibhavatlan

pyatlan/client/workflow.py

@@ -160,35 +160,44 @@ def find_run_by_id(self, id: str) -> Optional[WorkflowSearchResult]:
         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
+    def find_runs_by_status_and_time_range(
+        self,
+        status: List[AtlanWorkflowPhase],
+        started_at: Optional[str] = None,
+        finished_at: Optional[str] = None,
     ) -> List[WorkflowSearchResult]:
         """
-        Find workflows based on their status and interval
+        Find workflows by status and optional time filters on startedAt and/or finishedAt.
 
-        :param status: list of the status of the workflows to filter with
-        :param 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
+        :param status: list of the workflow statuses to filter
+        :param started_at: (optional) lower bound on 'status.startedAt' (e.g 'now-2h')
+        :param finished_at: (optional) lower bound on 'status.finishedAt' (e.g 'now-1h')
+        :returns: list of workflows matching the filters
+        :raises ValidationError: if inputs are invalid
         :raises AtlanError: on any API communication issue
         """
+        time_filters = []
+
+        if started_at:
+            time_filters.append(Range(field="status.startedAt", gte=started_at))
+        if finished_at:
+            time_filters.append(Range(field="status.finishedAt", gte=finished_at))
 
         run_lookup_query = Bool(
             must=[
                 NestedQuery(
                     query=Terms(
                         field="metadata.labels.workflows.argoproj.io/phase.keyword",
-                        values=[state.value for state in status],
+                        values=[s.value for s in status],
                     ),
                     path="metadata",
                 ),
-                Range(field="status.finishedAt", gt=f"now-{interval}h"),
+                *time_filters,
                 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)

tests/integration/test_workflow_client.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2024 Atlan Pte. Ltd.
 import time
+from datetime import datetime, timedelta, timezone
 from typing import Generator
 
 import pytest
@@ -171,18 +172,23 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
         and workflow_run.source.status.phase == AtlanWorkflowPhase.FAILED
     )
 
-
-def test_workflow_find_by_status_and_interval(client: AtlanClient):
-    runs_status = client.workflow.find_by_status_and_interval(
-        [AtlanWorkflowPhase.FAILED], 1
+    # Test find run by status and time range
+    runs_status = client.workflow.find_runs_by_status_and_time_range(
+        [AtlanWorkflowPhase.FAILED], started_at="now-1h"
     )
     assert runs_status
     workflow_run_status = runs_status[0]
+    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)
+    current_time = datetime.now(timezone.utc)
+    time_diff = current_time - start_datetime
     assert (
         workflow_run_status
         and workflow_run_status.source
         and workflow_run_status.source.status
         and workflow_run_status.source.status.phase == AtlanWorkflowPhase.FAILED
+        and time_diff < timedelta(hours=1)
     )
 
 

tests/unit/test_workflow_client.py

@@ -272,14 +272,17 @@ def test_find_by_type(client: WorkflowClient, mock_api_caller):
     )
 
 
-def test_find_by_status_and_interval(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}}}
     mock_api_caller._call_api.return_value = raw_json
 
     status = [AtlanWorkflowPhase.SUCCESS, AtlanWorkflowPhase.FAILED]
-    interval = 50
+    started_at = "now-2h"
+    finished_at = "now-1h"
 
-    assert client.find_by_status_and_interval(status, interval) == []
+    assert (
+        client.find_runs_by_status_and_time_range(status, started_at, finished_at) == []
+    )
     mock_api_caller._call_api.assert_called_once()
     assert isinstance(
         mock_api_caller._call_api.call_args.kwargs["request_obj"], WorkflowSearchRequest