"""Tests for src/queue_eta — slot-simulation ETA calculation.""" import json import math import os import tempfile import pytest from src.queue_eta import ( _LARGE_MODEL_HOURS, _SIZE_THRESHOLD_B, _SMALL_MODEL_HOURS, compute_single_eta, estimate_task_hours, format_eta, _get_params, ) # ═══════════════════════════════════════════════════════════════════════ # Helpers # ═══════════════════════════════════════════════════════════════════════ def _write(directory, owner, filename, data): d = os.path.join(directory, owner) os.makedirs(d, exist_ok=True) fpath = os.path.join(d, filename) with open(fpath, "w") as f: json.dump(data, f) return fpath # ═══════════════════════════════════════════════════════════════════════ # estimate_task_hours # ═══════════════════════════════════════════════════════════════════════ class TestEstimateTaskHours: def test_small_model(self): assert estimate_task_hours(7.0) == _SMALL_MODEL_HOURS def test_at_threshold(self): assert estimate_task_hours(_SIZE_THRESHOLD_B) == _SMALL_MODEL_HOURS def test_large_model(self): assert estimate_task_hours(70.0) == _LARGE_MODEL_HOURS def test_none_defaults_small(self): assert estimate_task_hours(None) == _SMALL_MODEL_HOURS def test_zero_defaults_small(self): assert estimate_task_hours(0) == _SMALL_MODEL_HOURS def test_negative_defaults_small(self): assert estimate_task_hours(-5) == _SMALL_MODEL_HOURS # ═══════════════════════════════════════════════════════════════════════ # format_eta # ═══════════════════════════════════════════════════════════════════════ class TestFormatEta: def test_zero(self): assert format_eta(0) == "< 1h" def test_negative(self): assert format_eta(-1) == "< 1h" def test_fraction(self): assert format_eta(0.5) == "~1h" def test_exact_hours(self): assert format_eta(6.0) == "~6h" def test_fractional_hours(self): assert format_eta(6.3) == "~7h" def test_exact_day(self): assert format_eta(24.0) == "~1d" def test_day_and_hours(self): assert format_eta(27.0) == "~1d 3h" def test_multi_days(self): assert format_eta(50.0) == "~2d 2h" # ═══════════════════════════════════════════════════════════════════════ # _get_params # ═══════════════════════════════════════════════════════════════════════ class TestGetParams: def test_model_params(self): assert _get_params({"model_params": 7.0}) == 7.0 def test_params_fallback(self): assert _get_params({"params": 14}) == 14.0 def test_model_params_preferred(self): assert _get_params({"model_params": 7, "params": 14}) == 7.0 def test_none_when_missing(self): assert _get_params({}) is None def test_string_value(self): assert _get_params({"model_params": "7.5"}) == 7.5 def test_bad_string(self): assert _get_params({"model_params": "unknown"}) is None # ═══════════════════════════════════════════════════════════════════════ # compute_single_eta — used by submit.py # ═══════════════════════════════════════════════════════════════════════ class TestComputeSingleEta: def setup_method(self): self._tmpdir = tempfile.mkdtemp(prefix="test_single_eta_") def test_empty_queue(self): """No running, no pending — new model is position 1.""" eta = compute_single_eta(self._tmpdir, model_params=7.0, concurrency=2) # running_remaining = 0, queue_pos = 1 # 0 + ceil(1/2)*3 = 3 assert eta == 3.0 def test_with_2r_1p(self): """2 running + 1 pending → new model at pos 2.""" _write(self._tmpdir, "a", "r1.json", { "model": "a/r1", "status": "Running", "script": "auto_quant", "model_params": 7.0, }) _write(self._tmpdir, "b", "r2.json", { "model": "b/r2", "status": "Running", "script": "auto_eval", "model_params": 7.0, }) _write(self._tmpdir, "c", "p1.json", { "model": "c/p1", "status": "Pending", "script": "auto_quant", "model_params": 7.0, "submitted_time": "2026-04-01T00:00:00Z", }) eta = compute_single_eta(self._tmpdir, model_params=7.0, concurrency=2) # running_remaining = avg(3,3) = 3, queue_pos = 2 # 3 + ceil(2/2)*3 = 3+3 = 6 assert eta == 6.0 def test_large_model(self): eta = compute_single_eta(self._tmpdir, model_params=70.0, concurrency=2) # running_remaining = 0, queue_pos = 1 # 0 + ceil(1/2)*5 = 5 assert eta == 5.0 def test_nonexistent_dir(self): eta = compute_single_eta("/tmp/nonexistent_eta_dir", model_params=7.0, concurrency=2) assert eta == 3.0