Fix tests

Author: r4victor

src/dstack/_internal/server/testing/common.py

@@ -302,6 +302,7 @@ async def create_job(
     job_num: int = 0,
     replica_num: int = 0,
     instance_assigned: bool = False,
+    disconnected_at: Optional[datetime] = None,
 ) -> JobModel:
     run_spec = RunSpec.parse_raw(run.run_spec)
     job_spec = (await get_job_specs_from_run_spec(run_spec, replica_num=replica_num))[0]
@@ -323,6 +324,7 @@ async def create_job(
         instance=instance,
         instance_assigned=instance_assigned,
         used_instance_id=instance.id if instance is not None else None,
+        disconnected_at=disconnected_at,
     )
     session.add(job)
     await session.commit()

src/tests/_internal/server/background/tasks/test_process_running_jobs.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from typing import Optional
 from unittest.mock import MagicMock, Mock, patch
@@ -490,6 +490,17 @@ async def test_pulling_shim_failed(self, test_db, session: AsyncSession):
             assert SSHTunnelMock.call_count == 3
         await session.refresh(job)
         assert job is not None
+        assert job.disconnected_at is not None
+        assert job.status == JobStatus.PULLING
+        with (
+            patch("dstack._internal.server.services.runner.ssh.SSHTunnel") as SSHTunnelMock,
+            patch("dstack._internal.server.services.runner.ssh.time.sleep"),
+            freeze_time(job.disconnected_at + timedelta(minutes=5)),
+        ):
+            SSHTunnelMock.side_effect = SSHError
+            await process_running_jobs()
+            assert SSHTunnelMock.call_count == 3
+        await session.refresh(job)
         assert job.status == JobStatus.TERMINATING
         assert job.termination_reason == JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY
         assert job.remove_at is None