tests/test_queue_eta.py

"""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