File size: 1,866 Bytes
77324b8 | 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 | from types import SimpleNamespace
import pytest
from agent.core import cost_estimation
def test_parse_timeout_hours_common_units():
assert cost_estimation.parse_timeout_hours(None) == 0.5
assert cost_estimation.parse_timeout_hours("30m") == 0.5
assert cost_estimation.parse_timeout_hours("3h") == 3
assert cost_estimation.parse_timeout_hours(3600) == 1
assert cost_estimation.parse_timeout_hours("not-a-duration") is None
@pytest.mark.asyncio
async def test_estimate_hf_job_cost_uses_catalog_price(monkeypatch):
async def fake_catalog():
return {"a100-large": 4.0}
monkeypatch.setattr(cost_estimation, "hf_jobs_price_catalog", fake_catalog)
estimate = await cost_estimation.estimate_hf_job_cost(
{"hardware_flavor": "a100-large", "timeout": "8h"}
)
assert estimate.estimated_cost_usd == 32.0
assert estimate.billable is True
@pytest.mark.asyncio
async def test_estimate_hf_job_cost_blocks_unknown_price(monkeypatch):
async def fake_catalog():
return {}
monkeypatch.setattr(cost_estimation, "hf_jobs_price_catalog", fake_catalog)
estimate = await cost_estimation.estimate_hf_job_cost(
{"hardware_flavor": "mystery-gpu", "timeout": "30m"}
)
assert estimate.estimated_cost_usd is None
assert estimate.billable is True
assert "No price" in estimate.block_reason
@pytest.mark.asyncio
async def test_estimate_sandbox_cost_is_zero_for_existing_or_cpu_basic():
existing = await cost_estimation.estimate_sandbox_cost(
{"hardware": "a100-large"},
session=SimpleNamespace(sandbox=object()),
)
cpu = await cost_estimation.estimate_sandbox_cost({"hardware": "cpu-basic"})
assert existing.estimated_cost_usd == 0.0
assert existing.billable is False
assert cpu.estimated_cost_usd == 0.0
assert cpu.billable is False
|