chore(python-sdk): Update tests

Author: nahua-aignx

tests/aignostics/application/utils_test.py

@@ -229,62 +229,133 @@ def test_application_run_status_to_str_terminated() -> None:
 # Tests for is_not_terminated_with_deadline_exceeded
 
 
+# --- Tests using scheduling response object (new primary path) ---
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_object_terminated_run() -> None:
+    """Test that terminated runs always return None regardless of scheduling deadline."""
+    past_deadline = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
+    scheduling = Mock(deadline=past_deadline)
+    result = is_not_terminated_with_deadline_exceeded(RunState.TERMINATED, scheduling=scheduling)
+    assert result is None
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_object_pending_future() -> None:
+    """Test that a pending run with scheduling deadline in the future returns False."""
+    from datetime import timedelta
+
+    future_deadline = datetime.now(tz=UTC) + timedelta(hours=1)
+    scheduling = Mock(deadline=future_deadline)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, scheduling=scheduling)
+    assert result is False
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_object_pending_past() -> None:
+    """Test that a pending run with scheduling deadline in the past returns True."""
+    past_deadline = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
+    scheduling = Mock(deadline=past_deadline)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, scheduling=scheduling)
+    assert result is True
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_object_processing_past() -> None:
+    """Test that a processing run with scheduling deadline in the past returns True."""
+    past_deadline = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
+    scheduling = Mock(deadline=past_deadline)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PROCESSING, scheduling=scheduling)
+    assert result is True
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_object_none_deadline() -> None:
+    """Test that scheduling object with None deadline returns None."""
+    scheduling = Mock(deadline=None)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, scheduling=scheduling)
+    assert result is None
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_none() -> None:
+    """Test that None scheduling and None metadata returns None."""
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, scheduling=None, custom_metadata=None)
+    assert result is None
+
+
+@pytest.mark.unit
+def test_is_not_terminated_with_deadline_exceeded_scheduling_takes_precedence() -> None:
+    """Test that scheduling object takes precedence over custom_metadata."""
+    from datetime import timedelta
+
+    # scheduling says future (not exceeded), custom_metadata says past (exceeded)
+    future_deadline = datetime.now(tz=UTC) + timedelta(hours=1)
+    scheduling = Mock(deadline=future_deadline)
+    metadata = {"sdk": {"scheduling": {"deadline": "2020-01-01T12:00:00Z"}}}
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, scheduling=scheduling, custom_metadata=metadata)
+    assert result is False  # scheduling wins
+
+
+# --- Tests using custom_metadata fallback (backward compatibility for older runs) ---
+
+
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_terminated_run() -> None:
     """Test that terminated runs always return None regardless of deadline."""
     past_deadline = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
     metadata = {"sdk": {"scheduling": {"deadline": past_deadline.isoformat()}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.TERMINATED, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.TERMINATED, custom_metadata=metadata)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_none_metadata() -> None:
     """Test that None metadata returns None."""
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, None)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=None)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_empty_metadata() -> None:
     """Test that empty metadata returns None."""
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, {})
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata={})
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_no_sdk_key() -> None:
     """Test that metadata without 'sdk' key returns None."""
     metadata = {"other_key": "other_value"}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_no_scheduling_key() -> None:
     """Test that metadata without 'scheduling' key returns None."""
     metadata = {"sdk": {"other_key": "other_value"}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_no_deadline_key() -> None:
     """Test that metadata without 'deadline' key returns None."""
     metadata = {"sdk": {"scheduling": {"other_key": "other_value"}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_pending_deadline_in_future() -> None:
     """Test that a pending run with deadline in the future returns False."""
-    # Create a deadline 1 hour in the future
     from datetime import timedelta
 
     future_deadline = datetime.now(tz=UTC) + timedelta(hours=1)
     metadata = {"sdk": {"scheduling": {"deadline": future_deadline.isoformat()}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is False
 
 
@@ -293,7 +364,7 @@ def test_is_not_terminated_with_deadline_exceeded_pending_deadline_in_past() ->
     """Test that a pending run with deadline in the past returns True."""
     past_deadline = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
     metadata = {"sdk": {"scheduling": {"deadline": past_deadline.isoformat()}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is True
 
 
@@ -302,15 +373,15 @@ def test_is_not_terminated_with_deadline_exceeded_processing_deadline_in_past()
     """Test that a processing run with deadline in the past returns True."""
     past_deadline = datetime(2020, 1, 1, 12, 0, 0, tzinfo=UTC)
     metadata = {"sdk": {"scheduling": {"deadline": past_deadline.isoformat()}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PROCESSING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PROCESSING, custom_metadata=metadata)
     assert result is True
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_invalid_datetime_format() -> None:
     """Test that invalid datetime format returns None."""
     metadata = {"sdk": {"scheduling": {"deadline": "not-a-valid-datetime"}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 
@@ -319,7 +390,7 @@ def test_is_not_terminated_with_deadline_exceeded_deadline_with_z_suffix() -> No
     """Test that deadline with Z suffix (UTC) is handled correctly for pending run."""
     past_deadline = "2020-01-01T12:00:00Z"
     metadata = {"sdk": {"scheduling": {"deadline": past_deadline}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is True
 
 
@@ -328,31 +399,31 @@ def test_is_not_terminated_with_deadline_exceeded_deadline_with_timezone_offset(
     """Test that deadline with timezone offset is handled correctly for pending run."""
     past_deadline = "2020-01-01T12:00:00+00:00"
     metadata = {"sdk": {"scheduling": {"deadline": past_deadline}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is True
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_deadline_empty_string() -> None:
     """Test that empty string deadline returns None."""
     metadata = {"sdk": {"scheduling": {"deadline": ""}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_deadline_none_value() -> None:
     """Test that None deadline value returns None."""
     metadata = {"sdk": {"scheduling": {"deadline": None}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 
 @pytest.mark.unit
 def test_is_not_terminated_with_deadline_exceeded_deadline_numeric_value() -> None:
     """Test that numeric deadline value returns None."""
     metadata = {"sdk": {"scheduling": {"deadline": 123456789}}}
-    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, metadata)
+    result = is_not_terminated_with_deadline_exceeded(RunState.PENDING, custom_metadata=metadata)
     assert result is None
 
 

tests/aignostics/platform/e2e_test.py

@@ -280,13 +280,13 @@ def _submit_and_validate(  # noqa: PLR0913, PLR0917
 
     logger.trace(f"Submitting application run for {application_id} version {application_version}")
     client = platform.Client()
+    scheduling = {
+        "due_date": (datetime.now(tz=UTC) + timedelta(seconds=due_date_seconds)).isoformat(),
+        "deadline": deadline.isoformat(),
+    }
     custom_metadata = {
         "sdk": {
             "tags": tags or set(),
-            "scheduling": {
-                "due_date": (datetime.now(tz=UTC) + timedelta(seconds=due_date_seconds)).isoformat(),
-                "deadline": deadline.isoformat(),
-            },
             "pipeline": {
                 "gpu": {
                     "gpu_type": PIPELINE_GPU_TYPE,
@@ -310,6 +310,7 @@ def _submit_and_validate(  # noqa: PLR0913, PLR0917
         application_version=application_version,
         items=payload,
         custom_metadata=custom_metadata,
+        scheduling=scheduling,
     )
 
     # Let's validate we can fiond the run by id
@@ -430,9 +431,20 @@ def _find_and_validate(
         sdk_metadata = RunSdkMetadata.model_validate(details.custom_metadata.get("sdk", {}))
         logger.trace(sdk_metadata.model_dump_json(indent=2))
         print(sdk_metadata.model_dump_json(indent=2))
-        allowed_duration = datetime.fromisoformat(sdk_metadata.scheduling.deadline) - datetime.fromisoformat(
-            sdk_metadata.submission.date
-        )
+        # Read deadline from the API response scheduling field (preferred),
+        # falling back to custom_metadata for runs submitted before this change
+        scheduling = getattr(details, "scheduling", None)
+        if scheduling is not None and getattr(scheduling, "deadline", None) is not None:
+            deadline_dt = (
+                scheduling.deadline
+                if isinstance(scheduling.deadline, datetime)
+                else datetime.fromisoformat(str(scheduling.deadline))
+            )
+            deadline_iso = deadline_dt.isoformat()
+        else:
+            deadline_iso = sdk_metadata.scheduling.deadline
+            deadline_dt = datetime.fromisoformat(deadline_iso)
+        allowed_duration = deadline_dt - datetime.fromisoformat(sdk_metadata.submission.date)
         allowed_hours = round(allowed_duration.total_seconds() / (60 * 60))
         deadline_met = details.state is RunState.TERMINATED
         metrics_run_attributes = {
@@ -441,7 +453,7 @@ def _find_and_validate(
             "application_version": application_version,
             "allowed_hours": allowed_hours,
             "submitted_at": sdk_metadata.submission.date,
-            "deadline": sdk_metadata.scheduling.deadline,
+            "deadline": deadline_iso,
             "state": details.state.value,
             "error_message": details.error_message,
             "error_code": details.error_code,