75 lines
3.0 KiB
Python
75 lines
3.0 KiB
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
from grm import stream # noqa: E402
|
|
|
|
|
|
def _ss(result=None):
|
|
ss = stream.StreamState()
|
|
if result is not None:
|
|
ss.feed(dict(result, type="result"), 1.0)
|
|
return ss
|
|
|
|
|
|
class ClassifyTest(unittest.TestCase):
|
|
def test_api_errors_map_to_reasons(self):
|
|
"""B1: an API error result gives status error with its mapped reason."""
|
|
cases = [({"is_error": True, "api_error_status": 429}, "usage_limit"),
|
|
({"is_error": True, "api_error_status": 529}, "overloaded"),
|
|
({"is_error": True, "result": "Prompt is too long"}, "prompt_too_long")]
|
|
for res, reason in cases:
|
|
self.assertEqual(stream.api_error_reason(res)[0], reason)
|
|
self.assertEqual(stream.classify(_ss(res), 1, False), ("error", reason))
|
|
|
|
def test_no_result(self):
|
|
"""B1: no result event is an error, by exit code or no_result."""
|
|
self.assertEqual(stream.classify(_ss(), 0, False), ("error", "no_result"))
|
|
self.assertEqual(stream.classify(_ss(), 1, False), ("error", "exit_code"))
|
|
self.assertEqual(stream.classify(_ss({"is_error": False}), 0, False), ("ok", None))
|
|
|
|
|
|
class UsageTest(unittest.TestCase):
|
|
def test_cache_miss(self):
|
|
"""B2: a first-request cache_read of 0 is a miss (reported, not an error)."""
|
|
self.assertEqual(stream.cache_verdict({"cache_read_input_tokens": 0}, 50000), "miss")
|
|
|
|
def test_two_models_and_cost_delta(self):
|
|
"""B5: both models in models_used; cost delta against the base, rounded at 6."""
|
|
before = {"claude-opus-5-5": {"inputTokens": 10, "outputTokens": 5}}
|
|
res = {"usage": {"input_tokens": 3, "output_tokens": 7},
|
|
"modelUsage": {"claude-opus-5-5": {"inputTokens": 12, "outputTokens": 10},
|
|
"claude-haiku-5-5": {"inputTokens": 1, "outputTokens": 1}},
|
|
"total_cost_usd": 0.12345678}
|
|
u = stream.usage_from_result(res, before, 0.1)
|
|
self.assertEqual(u["models_used"], ["claude-opus-5-5", "claude-haiku-5-5"])
|
|
self.assertEqual(u["cost_delta"], round(0.12345678 - 0.1, 6))
|
|
self.assertEqual(u["total"], 0.123457)
|
|
|
|
def test_missing_cost(self):
|
|
"""B5: no total_cost_usd gives cost_delta None."""
|
|
self.assertIsNone(stream.usage_from_result({"usage": {}}, {}, 0.5)["cost_delta"])
|
|
|
|
|
|
class WarmthTest(unittest.TestCase):
|
|
def test_no_hit_is_cold(self):
|
|
"""B6"""
|
|
self.assertEqual(stream.warmth(None, 3600, 1000.0)["base_state"], "cold")
|
|
|
|
def test_5m_only_ttl(self):
|
|
"""B6"""
|
|
self.assertEqual(stream.infer_ttl({"cache_creation_5m": 10, "cache_creation_1h": 0}, None), 300)
|
|
|
|
def test_expiring(self):
|
|
"""B6: a hit at now - ttl + 60 is expiring."""
|
|
now = 1_000_000.0
|
|
w = stream.warmth(now - 3600 + 60, 3600, now)
|
|
self.assertEqual(w["base_state"], "expiring")
|
|
self.assertEqual(w["ttl_left_s"], 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|