python(feat): Add better report creation and wait_until_complete to SiftClient

python/lib/sift_client/_internal/low_level_wrappers/reports.py

@@ -16,7 +16,7 @@
 from sift.reports.v1.reports_pb2_grpc import ReportServiceStub
 
 from sift_client._internal.low_level_wrappers.base import LowLevelClientBase
-from sift_client.sift_types.report import Report, ReportUpdate
+from sift_client.sift_types.report import PendingReport, Report, ReportUpdate
 from sift_client.transport import WithGrpcClient
 
 if TYPE_CHECKING:
@@ -130,14 +130,14 @@ async def list_all_reports(
             page_size=page_size,
         )
 
-    async def rerun_report(self, report_id: str) -> tuple[str, str]:
+    async def rerun_report(self, report_id: str) -> PendingReport:
         """Rerun a report.
 
         Args:
             report_id: The ID of the report to rerun.
 
         Returns:
-            A tuple of (job_id, new_report_id).
+            A PendingReport for the new report run.
 
         Raises:
             ValueError: If report_id is not provided.
@@ -148,7 +148,7 @@ async def rerun_report(self, report_id: str) -> tuple[str, str]:
         request = RerunReportRequest(report_id=report_id)
         response = await self._grpc_client.get_stub(ReportServiceStub).RerunReport(request)
         response = cast("RerunReportResponse", response)
-        return response.job_id, response.report_id
+        return PendingReport(report_id=response.report_id, job_id=response.job_id)
 
     async def cancel_report(self, report_id: str) -> None:
         """Cancel a report.

python/lib/sift_client/_internal/low_level_wrappers/rules.py

@@ -38,9 +38,9 @@
 from sift.rules.v1.rules_pb2_grpc import RuleServiceStub
 
 from sift_client._internal.low_level_wrappers.base import DEFAULT_PAGE_SIZE, LowLevelClientBase
-from sift_client._internal.low_level_wrappers.reports import ReportsLowLevelClient
 from sift_client._internal.util.timestamp import to_pb_timestamp
 from sift_client._internal.util.util import count_non_none
+from sift_client.sift_types.report import PendingReport
 from sift_client.sift_types.rule import (
     Rule,
     RuleCreate,
@@ -53,7 +53,6 @@
     from datetime import datetime
 
     from sift_client.sift_types.channel import ChannelReference
-    from sift_client.sift_types.report import Report
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -520,7 +519,7 @@ async def evaluate_rules(
         report_name: str | None = None,
         tags: list[str | Tag] | None = None,
         organization_id: str | None = None,
-    ) -> tuple[int, Report | None, str | None]:
+    ) -> PendingReport:
         """Evaluate a rule.
 
         Args:
@@ -588,10 +587,4 @@ async def evaluate_rules(
             request
         )
         response = cast("EvaluateRulesResponse", response)
-        created_annotation_count = response.created_annotation_count
-        report_id = response.report_id
-        job_id = response.job_id
-        if report_id:
-            report = await ReportsLowLevelClient(self._grpc_client).get_report(report_id=report_id)
-            return created_annotation_count, report, job_id
-        return created_annotation_count, None, job_id
+        return PendingReport._from_proto(response)

python/lib/sift_client/_tests/resources/test_jobs.py

@@ -8,6 +8,7 @@
 """
 
 from datetime import datetime, timedelta, timezone
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from grpc.aio import AioRpcError
@@ -291,6 +292,106 @@ async def test_retry_finished_job_no_effect(self, jobs_api_async):
                 with pytest.raises(AioRpcError, match="job cannot be retried"):
                     await jobs_api_async.retry(job)
 
+    class TestWaitUntilComplete:
+        """Tests for the async wait_until_complete method."""
+
+        @pytest.mark.asyncio
+        async def test_returns_immediately_when_job_already_complete(self, jobs_api_async):
+            """When get returns a completed job on first call, wait returns immediately."""
+            job_id = "test-job-id"
+            mock_job = MagicMock()
+            mock_job.job_status = JobStatus.FINISHED
+
+            with patch(
+                "sift_client.resources.jobs.JobsAPIAsync.get",
+                new_callable=AsyncMock,
+                return_value=mock_job,
+            ) as mock_get:
+                result = await jobs_api_async.wait_until_complete(job=job_id)
+
+            assert result is mock_job
+            assert result.job_status == JobStatus.FINISHED
+            mock_get.assert_called_once_with(job_id)
+
+        @pytest.mark.asyncio
+        async def test_returns_immediately_when_job_already_failed(self, jobs_api_async):
+            """When get returns a failed job on first call, wait returns immediately."""
+            job_id = "test-job-id"
+            mock_job = MagicMock()
+            mock_job.job_status = JobStatus.FAILED
+
+            with patch(
+                "sift_client.resources.jobs.JobsAPIAsync.get",
+                new_callable=AsyncMock,
+                return_value=mock_job,
+            ) as mock_get:
+                result = await jobs_api_async.wait_until_complete(job=job_id)
+
+            assert result is mock_job
+            assert result.job_status == JobStatus.FAILED
+            mock_get.assert_called_once_with(job_id)
+
+        @pytest.mark.asyncio
+        async def test_returns_immediately_when_job_already_cancelled(self, jobs_api_async):
+            """When get returns a cancelled job on first call, wait returns immediately."""
+            job_id = "test-job-id"
+            mock_job = MagicMock()
+            mock_job.job_status = JobStatus.CANCELLED
+
+            with patch(
+                "sift_client.resources.jobs.JobsAPIAsync.get",
+                new_callable=AsyncMock,
+                return_value=mock_job,
+            ) as mock_get:
+                result = await jobs_api_async.wait_until_complete(job=job_id)
+
+            assert result is mock_job
+            assert result.job_status == JobStatus.CANCELLED
+            mock_get.assert_called_once_with(job_id)
+
+        @pytest.mark.asyncio
+        async def test_polls_until_complete(self, jobs_api_async):
+            """When get returns running then finished, wait returns after second poll."""
+            job_id = "test-job-id"
+            running_job = MagicMock()
+            running_job.job_status = JobStatus.RUNNING
+            finished_job = MagicMock()
+            finished_job.job_status = JobStatus.FINISHED
+
+            with patch(
+                "sift_client.resources.jobs.JobsAPIAsync.get",
+                new_callable=AsyncMock,
+                side_effect=[running_job, finished_job],
+            ) as mock_get:
+                result = await jobs_api_async.wait_until_complete(
+                    job=job_id,
+                    polling_interval_secs=0.01,
+                    timeout_secs=10.0,
+                )
+
+            assert result is finished_job
+            assert result.job_status == JobStatus.FINISHED
+            assert mock_get.call_count == 2
+
+        @pytest.mark.asyncio
+        async def test_raises_timeout_error_when_not_complete_in_time(self, jobs_api_async):
+            """When job never reaches a completed state, TimeoutError is raised."""
+            job_id = "test-job-id"
+            running_job = MagicMock()
+            running_job.job_status = JobStatus.RUNNING
+
+            with patch(
+                "sift_client.resources.jobs.JobsAPIAsync.get",
+                new_callable=AsyncMock,
+                return_value=running_job,
+            ):
+                with pytest.raises(TimeoutError):
+                    await jobs_api_async.wait_until_complete(
+                        job=job_id,
+                        polling_interval_secs=0.05,
+                        timeout_secs=0.1,
+                    )
+
     class TestJobProperties:
         """Tests for job property methods."""
 

python/lib/sift_client/_tests/resources/test_reports.py

@@ -55,27 +55,61 @@ def test_client_binding(sift_client):
 @pytest.mark.integration
 class TestReports:
     def test_create_from_rules(self, nostromo_run, test_rule, sift_client):
-        report_from_rules = sift_client.reports.create_from_rules(
+        pending = sift_client.reports.create_from_rules(
             name="report_from_rules",
             run=nostromo_run,
             rules=[test_rule],
         )
-        assert report_from_rules is not None
-        assert report_from_rules.run_id == nostromo_run.id_
+        assert pending is not None
+        report = sift_client.reports.get(report_id=pending.report_id)
+        assert report.run_id == nostromo_run.id_
+
+    @pytest.mark.asyncio
+    async def test_wait_until_complete(self, nostromo_run, test_rule, sift_client):
+        """Create a report and wait for its job to complete via async wait_until_complete."""
+        pending = sift_client.reports.create_from_rules(
+            name="report_wait_until_complete",
+            run=nostromo_run,
+            rules=[test_rule],
+        )
+        assert pending is not None
+        assert pending.job_id
+
+        completed_report = await sift_client.async_.reports.wait_until_complete(
+            report=pending,
+            polling_interval_secs=2,
+            timeout_secs=120,
+        )
+
+        assert completed_report is not None
+        assert completed_report.id_ == pending.report_id
+        assert completed_report.job_id == pending.job_id
+
+        completed_rule_statuses = (
+            ReportRuleStatus.FINISHED,
+            ReportRuleStatus.FAILED,
+            ReportRuleStatus.CANCELED,
+            ReportRuleStatus.ERROR,
+        )
+        assert len(completed_report.summaries) == 1
+        assert any(s.status in completed_rule_statuses for s in completed_report.summaries), (
+            "expected rule summary to be in a completed state"
+        )
 
     def test_create_from_applicable_rules(
         self, test_rule, nostromo_asset, nostromo_run, sift_client
     ):
         if not test_rule.asset_ids:
             # Test rule may exist but be in a state where it no longer applies to the asset associated w/ the run so re-attach it if necessary.
             test_rule = test_rule.update(update={"asset_ids": [nostromo_asset._id_or_error]})
-        report_from_applicable_rules = sift_client.reports.create_from_applicable_rules(
+        pending = sift_client.reports.create_from_applicable_rules(
             name="report_from_applicable_rules_run",
             run=nostromo_run,
             organization_id=nostromo_run.organization_id,
         )
-        assert report_from_applicable_rules is not None
-        assert report_from_applicable_rules.run_id == nostromo_run.id_
+        assert pending is not None
+        report = sift_client.reports.get(report_id=pending.report_id)
+        assert report.run_id == nostromo_run.id_
 
     def test_list(self, nostromo_asset, nostromo_run, tags, sift_client):
         reports = sift_client.reports.list_(
@@ -85,27 +119,31 @@ def test_list(self, nostromo_asset, nostromo_run, tags, sift_client):
         assert len(reports) > 0
 
     def test_rerun(self, nostromo_asset, nostromo_run, test_rule, sift_client):
-        report_from_rules = sift_client.reports.create_from_rules(
+        pending = sift_client.reports.create_from_rules(
             name="report_from_rules",
             run=nostromo_run,
             rules=[test_rule],
         )
-        assert report_from_rules is not None
-        job_id, rerun_report_id = sift_client.reports.rerun(report=report_from_rules)
-        rerun_report = sift_client.reports.get(report_id=rerun_report_id)
+        assert pending is not None
+        rerun_pending = sift_client.reports.rerun(report=pending)
+        assert rerun_pending is not None
+        assert rerun_pending.report_id
+        assert rerun_pending.job_id
+        rerun_report = sift_client.reports.get(report_id=rerun_pending.report_id)
         assert rerun_report is not None
         assert rerun_report.run_id == nostromo_run.id_
-        assert rerun_report.rerun_from_report_id == report_from_rules.id_
+        assert rerun_report.rerun_from_report_id == pending.report_id
 
     def test_update(self, nostromo_asset, nostromo_run, test_rule, sift_client):
-        report_from_rules = sift_client.reports.create_from_rules(
+        pending = sift_client.reports.create_from_rules(
             name="report_from_rules",
             run=nostromo_run,
             rules=[test_rule],
         )
-        assert report_from_rules is not None
+        assert pending is not None
+        report = sift_client.reports.get(report_id=pending.report_id)
         updated_report = sift_client.reports.update(
-            report=report_from_rules,
+            report=report,
             update={
                 "metadata": {
                     "test_type": "ci",
@@ -120,43 +158,44 @@ def test_find_multiple(self, sift_client):
             sift_client.reports.find(name="report_from_rules")
 
     def test_cancel(self, nostromo_asset, nostromo_run, test_rule, sift_client):
-        report_from_rules = sift_client.reports.create_from_rules(
+        pending = sift_client.reports.create_from_rules(
             name="report_from_rules",
             run=nostromo_run,
             rules=[test_rule],
         )
-        assert report_from_rules is not None
-        job_id, second_rerun_report_id = sift_client.reports.rerun(report=report_from_rules)
-        assert second_rerun_report_id is not None
-        sift_client.reports.cancel(report=second_rerun_report_id)
-        canceled_report = sift_client.reports.find(report_ids=[second_rerun_report_id])
+        assert pending is not None
+        second_rerun_pending = sift_client.reports.rerun(report=pending)
+        assert second_rerun_pending is not None
+        sift_client.reports.cancel(report=second_rerun_pending)
+        canceled_report = sift_client.reports.find(report_ids=[second_rerun_pending.report_id])
         assert canceled_report is not None
         for summary in canceled_report.summaries:
             # Sometimes the report finishes before it can be canceled.
             assert summary.status in [ReportRuleStatus.CANCELED, ReportRuleStatus.FINISHED]
 
     def test_archive(self, nostromo_run, test_rule, sift_client):
-        report_from_rules = sift_client.reports.create_from_rules(
+        pending = sift_client.reports.create_from_rules(
             name="report_from_rules",
             run=nostromo_run,
             rules=[test_rule],
         )
-        assert report_from_rules is not None
-        archived_report = sift_client.reports.archive(report=report_from_rules)
+        assert pending is not None
+        report = pending.wait_until_complete(polling_interval_secs=2, timeout_secs=120)
+        archived_report = sift_client.reports.archive(report=report)
         assert archived_report is not None
         assert archived_report.is_archived == True
 
-    def test_unarchive(self, sift_client):
-        reports_from_rules = sift_client.reports.list_(
-            name="report_from_rules", include_archived=True
+    def test_unarchive(self, nostromo_run, test_rule, sift_client):
+        # Create, wait for completion, then archive to ensure we have an archived report
+        pending = sift_client.reports.create_from_rules(
+            name="report_from_rules_unarchive",
+            run=nostromo_run,
+            rules=[test_rule],
         )
-        report_from_rules = None
-        for report_from_rules in reports_from_rules:
-            if report_from_rules.is_archived:
-                report_from_rules = report_from_rules
-                break
-        assert report_from_rules is not None
-        assert report_from_rules.is_archived == True
-        unarchived_report = sift_client.reports.unarchive(report=report_from_rules)
+        assert pending is not None
+        report = sift_client.reports.get(report_id=pending.report_id)
+        archived_report = sift_client.reports.archive(report=report)
+        assert archived_report.is_archived is True
+        unarchived_report = sift_client.reports.unarchive(report=archived_report)
         assert unarchived_report is not None
-        assert unarchived_report.is_archived == False
+        assert unarchived_report.is_archived is False