Spaces:
Running
on
T4
Running
on
T4
File size: 10,971 Bytes
722753e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 |
"""Unit tests for the async job store.
Tests the JobStore class that manages background ML inference jobs.
Follows Uncle Bob's testing principles:
- Test behavior, not implementation
- Each test verifies one thing
- Tests are independent and repeatable
"""
from collections.abc import Generator
from datetime import datetime, timedelta
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import patch
import pytest
from stroke_deepisles_demo.api.job_store import (
Job,
JobStatus,
JobStore,
get_job_store,
init_job_store,
)
class TestJob:
"""Tests for the Job dataclass."""
def test_new_job_has_zero_elapsed_seconds(self) -> None:
"""A job that hasn't started should report 0 elapsed seconds."""
job = Job(
id="abc123",
status=JobStatus.PENDING,
case_id="sub-stroke0001",
fast_mode=True,
created_at=datetime.now(),
)
assert job.elapsed_seconds == 0.0
def test_running_job_tracks_elapsed_time(self) -> None:
"""A running job should report elapsed time since start."""
start = datetime.now() - timedelta(seconds=10)
job = Job(
id="abc123",
status=JobStatus.RUNNING,
case_id="sub-stroke0001",
fast_mode=True,
created_at=start - timedelta(seconds=1),
started_at=start,
)
# Should be approximately 10 seconds (with some tolerance)
assert 9.5 <= job.elapsed_seconds <= 11.0
def test_completed_job_has_fixed_elapsed_time(self) -> None:
"""A completed job should report time from start to completion."""
start = datetime.now() - timedelta(seconds=30)
end = start + timedelta(seconds=15)
job = Job(
id="abc123",
status=JobStatus.COMPLETED,
case_id="sub-stroke0001",
fast_mode=True,
created_at=start - timedelta(seconds=1),
started_at=start,
completed_at=end,
)
# Should be exactly 15 seconds (completed job doesn't change)
assert job.elapsed_seconds == 15.0
def test_to_dict_includes_required_fields(self) -> None:
"""Job.to_dict() should include all fields needed by the API."""
job = Job(
id="abc123",
status=JobStatus.RUNNING,
case_id="sub-stroke0001",
fast_mode=True,
created_at=datetime.now(),
started_at=datetime.now(),
progress=50,
progress_message="Processing...",
)
data = job.to_dict()
assert data["jobId"] == "abc123"
assert data["status"] == "running"
assert data["progress"] == 50
assert data["progressMessage"] == "Processing..."
assert "elapsedSeconds" in data
def test_to_dict_includes_result_when_completed(self) -> None:
"""Completed jobs should include result data in to_dict()."""
job = Job(
id="abc123",
status=JobStatus.COMPLETED,
case_id="sub-stroke0001",
fast_mode=True,
created_at=datetime.now(),
started_at=datetime.now(),
completed_at=datetime.now(),
result={"caseId": "sub-stroke0001", "diceScore": 0.847},
)
data = job.to_dict()
assert "result" in data
assert data["result"]["diceScore"] == 0.847
def test_to_dict_includes_error_when_failed(self) -> None:
"""Failed jobs should include error message in to_dict()."""
job = Job(
id="abc123",
status=JobStatus.FAILED,
case_id="sub-stroke0001",
fast_mode=True,
created_at=datetime.now(),
error="GPU out of memory",
)
data = job.to_dict()
assert "error" in data
assert data["error"] == "GPU out of memory"
class TestJobStore:
"""Tests for the JobStore class."""
@pytest.fixture
def store(self) -> Generator[JobStore, None, None]:
"""Create a fresh JobStore for each test."""
with TemporaryDirectory() as tmpdir:
yield JobStore(results_dir=Path(tmpdir))
def test_create_job_returns_pending_job(self, store: JobStore) -> None:
"""Creating a job should return a job in PENDING status."""
job = store.create_job("job-1", "sub-stroke0001", fast_mode=True)
assert job.id == "job-1"
assert job.status == JobStatus.PENDING
assert job.case_id == "sub-stroke0001"
assert job.fast_mode is True
def test_get_job_returns_created_job(self, store: JobStore) -> None:
"""get_job() should return a previously created job."""
store.create_job("job-1", "sub-stroke0001", fast_mode=True)
job = store.get_job("job-1")
assert job is not None
assert job.id == "job-1"
def test_get_job_returns_none_for_unknown_id(self, store: JobStore) -> None:
"""get_job() should return None for unknown job IDs."""
job = store.get_job("nonexistent")
assert job is None
def test_start_job_changes_status_to_running(self, store: JobStore) -> None:
"""start_job() should update job status to RUNNING."""
store.create_job("job-1", "sub-stroke0001", fast_mode=True)
store.start_job("job-1")
job = store.get_job("job-1")
assert job is not None
assert job.status == JobStatus.RUNNING
assert job.started_at is not None
def test_update_progress_changes_progress_fields(self, store: JobStore) -> None:
"""update_progress() should update progress and message."""
store.create_job("job-1", "sub-stroke0001", fast_mode=True)
store.start_job("job-1")
store.update_progress("job-1", 75, "Computing metrics...")
job = store.get_job("job-1")
assert job is not None
assert job.progress == 75
assert job.progress_message == "Computing metrics..."
def test_update_progress_clamps_to_valid_range(self, store: JobStore) -> None:
"""update_progress() should clamp progress to 0-100."""
store.create_job("job-1", "sub-stroke0001", fast_mode=True)
store.start_job("job-1")
store.update_progress("job-1", 150, "Over 100")
job = store.get_job("job-1")
assert job is not None
assert job.progress == 100
store.update_progress("job-1", -10, "Negative")
job = store.get_job("job-1")
assert job is not None
assert job.progress == 0
def test_complete_job_sets_status_and_result(self, store: JobStore) -> None:
"""complete_job() should mark job as completed with result."""
store.create_job("job-1", "sub-stroke0001", fast_mode=True)
store.start_job("job-1")
result = {"caseId": "sub-stroke0001", "diceScore": 0.847}
store.complete_job("job-1", result)
job = store.get_job("job-1")
assert job is not None
assert job.status == JobStatus.COMPLETED
assert job.progress == 100
assert job.result == result
assert job.completed_at is not None
def test_fail_job_sets_status_and_error(self, store: JobStore) -> None:
"""fail_job() should mark job as failed with error message."""
store.create_job("job-1", "sub-stroke0001", fast_mode=True)
store.start_job("job-1")
store.fail_job("job-1", "GPU out of memory")
job = store.get_job("job-1")
assert job is not None
assert job.status == JobStatus.FAILED
assert job.error == "GPU out of memory"
assert job.completed_at is not None
def test_len_returns_number_of_jobs(self, store: JobStore) -> None:
"""len(store) should return the number of jobs."""
assert len(store) == 0
store.create_job("job-1", "case1", fast_mode=True)
assert len(store) == 1
store.create_job("job-2", "case2", fast_mode=True)
assert len(store) == 2
class TestJobStoreCleanup:
"""Tests for job cleanup functionality."""
def test_cleanup_removes_old_completed_jobs(self) -> None:
"""cleanup_old_jobs() should remove jobs older than TTL."""
with TemporaryDirectory() as tmpdir:
# Use a very short TTL for testing
store = JobStore(ttl=timedelta(seconds=0), results_dir=Path(tmpdir))
store.create_job("job-1", "case1", fast_mode=True)
store.start_job("job-1")
store.complete_job("job-1", {"result": "data"})
# Job is "old" immediately (TTL=0)
cleaned = store.cleanup_old_jobs()
assert cleaned == 1
assert store.get_job("job-1") is None
def test_cleanup_keeps_running_jobs(self) -> None:
"""cleanup_old_jobs() should not remove running jobs."""
with TemporaryDirectory() as tmpdir:
store = JobStore(ttl=timedelta(seconds=0), results_dir=Path(tmpdir))
store.create_job("job-1", "case1", fast_mode=True)
store.start_job("job-1")
# Job is running, not completed
cleaned = store.cleanup_old_jobs()
assert cleaned == 0
assert store.get_job("job-1") is not None
def test_cleanup_removes_result_files(self) -> None:
"""cleanup_old_jobs() should also remove result files on disk."""
with TemporaryDirectory() as tmpdir:
results_dir = Path(tmpdir)
store = JobStore(ttl=timedelta(seconds=0), results_dir=results_dir)
# Create job and its result directory
store.create_job("job-1", "case1", fast_mode=True)
store.start_job("job-1")
job_results = results_dir / "job-1"
job_results.mkdir()
(job_results / "prediction.nii.gz").touch()
store.complete_job("job-1", {"result": "data"})
# Cleanup should remove both job record and files
store.cleanup_old_jobs()
assert not job_results.exists()
class TestGlobalJobStore:
"""Tests for the global job store singleton."""
def test_get_job_store_raises_before_init(self) -> None:
"""get_job_store() should raise if not initialized."""
# Patch the global to simulate uninitialized state
with (
patch("stroke_deepisles_demo.api.job_store.job_store", None),
pytest.raises(RuntimeError, match="not initialized"),
):
get_job_store()
def test_init_job_store_creates_global_instance(self) -> None:
"""init_job_store() should create and return a JobStore."""
with TemporaryDirectory() as tmpdir:
store = init_job_store(results_dir=Path(tmpdir))
assert store is not None
assert isinstance(store, JobStore)
# Clean up the scheduler
store.stop_cleanup_scheduler()
|