276 lines
9.9 KiB
Python
276 lines
9.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Tests for the household Grok usage collector.
|
||
|
|
|
||
|
|
No live network. No real credentials. Fixtures use obviously fake tokens.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import stat
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
from unittest import mock
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[1]
|
||
|
|
sys.path.insert(0, str(ROOT / "src"))
|
||
|
|
|
||
|
|
import omarchy_grok_usage as m # noqa: E402
|
||
|
|
|
||
|
|
FAKE_TOKEN = "fake-test-token-not-a-real-secret"
|
||
|
|
|
||
|
|
|
||
|
|
class TestAuthParse(unittest.TestCase):
|
||
|
|
def test_extracts_key_from_accounts_xai_entry(self) -> None:
|
||
|
|
payload = {
|
||
|
|
"https://accounts.x.ai/sign-in": {
|
||
|
|
"key": FAKE_TOKEN,
|
||
|
|
"first_name": "Test",
|
||
|
|
}
|
||
|
|
}
|
||
|
|
token, label = m.parse_auth(payload)
|
||
|
|
self.assertEqual(token, FAKE_TOKEN)
|
||
|
|
self.assertEqual(label, "Test")
|
||
|
|
|
||
|
|
def test_extracts_key_from_first_object_with_key(self) -> None:
|
||
|
|
payload = {
|
||
|
|
"other": {"nope": 1},
|
||
|
|
"grok-com": {"key": FAKE_TOKEN, "first_name": "Ada"},
|
||
|
|
}
|
||
|
|
token, label = m.parse_auth(payload)
|
||
|
|
self.assertEqual(token, FAKE_TOKEN)
|
||
|
|
self.assertEqual(label, "Ada")
|
||
|
|
|
||
|
|
def test_missing_key_is_unauthenticated(self) -> None:
|
||
|
|
token, label = m.parse_auth({"x": {"refresh": "nope"}})
|
||
|
|
self.assertIsNone(token)
|
||
|
|
self.assertIsNone(label)
|
||
|
|
|
||
|
|
def test_empty_key_is_unauthenticated(self) -> None:
|
||
|
|
token, _label = m.parse_auth({"https://accounts.x.ai/sign-in": {"key": ""}})
|
||
|
|
self.assertIsNone(token)
|
||
|
|
|
||
|
|
|
||
|
|
class TestBillingMap(unittest.TestCase):
|
||
|
|
def test_credit_usage_percent_and_period_end(self) -> None:
|
||
|
|
payload = {
|
||
|
|
"config": {
|
||
|
|
"creditUsagePercent": 31.5,
|
||
|
|
"currentPeriod": {"end": "2026-08-27T00:00:00Z"},
|
||
|
|
"monthlyLimit": {"val": 60000},
|
||
|
|
"used": {"val": 18900},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
limits = m.limits_from_billing(payload)
|
||
|
|
self.assertEqual(len(limits), 1)
|
||
|
|
self.assertEqual(limits[0]["label"], "Weekly")
|
||
|
|
self.assertEqual(limits[0]["percent"], 31.5)
|
||
|
|
self.assertEqual(limits[0]["resetsAt"], "2026-08-27T00:00:00+00:00")
|
||
|
|
|
||
|
|
def test_falls_back_to_used_over_limit(self) -> None:
|
||
|
|
payload = {
|
||
|
|
"config": {
|
||
|
|
"monthlyLimit": {"val": 100},
|
||
|
|
"used": {"val": 40},
|
||
|
|
"billingPeriodEnd": "2026-09-01T00:00:00+00:00",
|
||
|
|
}
|
||
|
|
}
|
||
|
|
limits = m.limits_from_billing(payload)
|
||
|
|
self.assertEqual(limits[0]["percent"], 40.0)
|
||
|
|
self.assertEqual(limits[0]["resetsAt"], "2026-09-01T00:00:00+00:00")
|
||
|
|
|
||
|
|
def test_empty_billing_yields_no_limits(self) -> None:
|
||
|
|
self.assertEqual(m.limits_from_billing({}), [])
|
||
|
|
|
||
|
|
def test_percent_clamped_to_100(self) -> None:
|
||
|
|
payload = {"config": {"creditUsagePercent": 140}}
|
||
|
|
limits = m.limits_from_billing(payload)
|
||
|
|
self.assertEqual(limits[0]["percent"], 100.0)
|
||
|
|
|
||
|
|
|
||
|
|
class TestTierLabel(unittest.TestCase):
|
||
|
|
def test_prefers_subscription_tier_display(self) -> None:
|
||
|
|
self.assertEqual(
|
||
|
|
m.tier_from_settings({"subscription_tier_display": "SuperGrok Heavy"}),
|
||
|
|
"SuperGrok Heavy",
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_splits_camel_subscription_tier(self) -> None:
|
||
|
|
self.assertEqual(
|
||
|
|
m.tier_from_settings({"subscriptionTier": "SuperGrokPro"}),
|
||
|
|
"SuperGrok Pro",
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_empty_settings(self) -> None:
|
||
|
|
self.assertEqual(m.tier_from_settings({}), "")
|
||
|
|
|
||
|
|
|
||
|
|
class TestOfficialYield(unittest.TestCase):
|
||
|
|
def test_yields_when_omarchy_path_collector_exists(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
official = Path(tmp) / "bin" / "omarchy-agent-usage-grok"
|
||
|
|
official.parent.mkdir(parents=True)
|
||
|
|
official.write_text("#!/bin/sh\n")
|
||
|
|
official.chmod(official.stat().st_mode | stat.S_IEXEC)
|
||
|
|
env = {"OMARCHY_PATH": tmp}
|
||
|
|
self.assertTrue(m.official_grok_collector_present(env=env))
|
||
|
|
|
||
|
|
def test_absent_when_no_official_binary(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
env = {
|
||
|
|
"OMARCHY_PATH": tmp,
|
||
|
|
"PATH": tmp,
|
||
|
|
}
|
||
|
|
self.assertFalse(
|
||
|
|
m.official_grok_collector_present(
|
||
|
|
env=env, extra_paths=(str(Path(tmp) / "nope"),)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_force_env_disables_yield(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
official = Path(tmp) / "bin" / "omarchy-agent-usage-grok"
|
||
|
|
official.parent.mkdir(parents=True)
|
||
|
|
official.write_text("#!/bin/sh\n")
|
||
|
|
official.chmod(official.stat().st_mode | stat.S_IEXEC)
|
||
|
|
env = {"OMARCHY_PATH": tmp, "OMARCHY_GROK_USAGE_FORCE": "1"}
|
||
|
|
self.assertFalse(m.official_grok_collector_present(env=env))
|
||
|
|
|
||
|
|
|
||
|
|
class TestRecord(unittest.TestCase):
|
||
|
|
def test_token_never_enters_record(self) -> None:
|
||
|
|
record = m.build_record(
|
||
|
|
limits=[{"label": "Weekly", "percent": 10.0, "resetsAt": None}],
|
||
|
|
tier_label="SuperGrok",
|
||
|
|
usage_status_text="",
|
||
|
|
auth_help_text="Run `grok login`.",
|
||
|
|
now_iso="2026-08-20T00:00:00+00:00",
|
||
|
|
)
|
||
|
|
blob = json.dumps(record)
|
||
|
|
self.assertNotIn(FAKE_TOKEN, blob)
|
||
|
|
self.assertNotIn("Authorization", blob)
|
||
|
|
self.assertEqual(record["id"], "grok")
|
||
|
|
self.assertEqual(record["schemaVersion"], 1)
|
||
|
|
self.assertTrue(record["ready"])
|
||
|
|
self.assertEqual(record["scope"], "account")
|
||
|
|
self.assertFalse(record["hasLocalStats"])
|
||
|
|
|
||
|
|
def test_unauthenticated_record_is_not_ready(self) -> None:
|
||
|
|
record = m.build_record(
|
||
|
|
limits=[],
|
||
|
|
tier_label="",
|
||
|
|
usage_status_text="",
|
||
|
|
auth_help_text="Run `grok login` to restore SuperGrok usage.",
|
||
|
|
now_iso="2026-08-20T00:00:00+00:00",
|
||
|
|
)
|
||
|
|
self.assertFalse(record["ready"])
|
||
|
|
self.assertEqual(record["limits"], [])
|
||
|
|
|
||
|
|
|
||
|
|
class TestWritePath(unittest.TestCase):
|
||
|
|
def test_skips_write_when_official_present(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
usage_dir = Path(tmp) / "usage"
|
||
|
|
usage_dir.mkdir()
|
||
|
|
existing = usage_dir / "grok.json"
|
||
|
|
existing.write_text('{"id":"grok","from":"official"}\n')
|
||
|
|
result = m.maybe_write_record(
|
||
|
|
record={"id": "grok", "from": "household"},
|
||
|
|
usage_dir=usage_dir,
|
||
|
|
yield_to_official=True,
|
||
|
|
)
|
||
|
|
self.assertEqual(result.action, "yielded")
|
||
|
|
self.assertEqual(existing.read_text(), '{"id":"grok","from":"official"}\n')
|
||
|
|
|
||
|
|
def test_writes_atomically_when_no_official(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
usage_dir = Path(tmp) / "usage"
|
||
|
|
result = m.maybe_write_record(
|
||
|
|
record={"id": "grok", "schemaVersion": 1},
|
||
|
|
usage_dir=usage_dir,
|
||
|
|
yield_to_official=False,
|
||
|
|
)
|
||
|
|
self.assertEqual(result.action, "wrote")
|
||
|
|
written = json.loads((usage_dir / "grok.json").read_text())
|
||
|
|
self.assertEqual(written["id"], "grok")
|
||
|
|
|
||
|
|
def test_https_only_client_rejects_http(self) -> None:
|
||
|
|
with self.assertRaises(m.HttpError):
|
||
|
|
m.https_get("http://example.com/billing", headers={})
|
||
|
|
|
||
|
|
|
||
|
|
class TestHttpsClientRedirects(unittest.TestCase):
|
||
|
|
def test_redirects_are_refused(self) -> None:
|
||
|
|
class FakeResponse:
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *args):
|
||
|
|
return False
|
||
|
|
|
||
|
|
def read(self) -> bytes:
|
||
|
|
return b"{}"
|
||
|
|
|
||
|
|
def geturl(self) -> str:
|
||
|
|
return "https://evil.example/steal"
|
||
|
|
|
||
|
|
def fake_open(_request, timeout=10):
|
||
|
|
return FakeResponse()
|
||
|
|
|
||
|
|
with mock.patch.object(m, "_urlopen", fake_open):
|
||
|
|
with self.assertRaises(m.HttpError) as ctx:
|
||
|
|
m.https_get(
|
||
|
|
"https://cli-chat-proxy.grok.com/v1/billing?format=credits",
|
||
|
|
headers={"Authorization": "Bearer " + FAKE_TOKEN},
|
||
|
|
)
|
||
|
|
self.assertIn("redirect", str(ctx.exception).lower())
|
||
|
|
|
||
|
|
|
||
|
|
class TestCollect(unittest.TestCase):
|
||
|
|
def test_collect_record_does_not_embed_token(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
grok = Path(tmp) / ".grok"
|
||
|
|
grok.mkdir()
|
||
|
|
(grok / "auth.json").write_text(
|
||
|
|
json.dumps({"https://accounts.x.ai/sign-in": {"key": FAKE_TOKEN}}),
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
|
||
|
|
def fake_get(url: str, headers):
|
||
|
|
auth = headers.get("Authorization", "")
|
||
|
|
self.assertTrue(auth.startswith("Bearer "))
|
||
|
|
if "billing" in url:
|
||
|
|
return {
|
||
|
|
"config": {
|
||
|
|
"creditUsagePercent": 12,
|
||
|
|
"currentPeriod": {"end": "2026-08-27T00:00:00Z"},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if "settings" in url:
|
||
|
|
return {"subscription_tier_display": "SuperGrok"}
|
||
|
|
raise AssertionError(url)
|
||
|
|
|
||
|
|
record = m.collect_record(
|
||
|
|
env={"GROK_HOME": str(grok)},
|
||
|
|
home=Path(tmp),
|
||
|
|
http_get=fake_get,
|
||
|
|
)
|
||
|
|
blob = json.dumps(record)
|
||
|
|
self.assertNotIn(FAKE_TOKEN, blob)
|
||
|
|
self.assertEqual(record["limits"][0]["percent"], 12)
|
||
|
|
self.assertEqual(record["tierLabel"], "SuperGrok")
|
||
|
|
self.assertTrue(record["ready"])
|
||
|
|
|
||
|
|
def test_collect_record_without_auth_file(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
record = m.collect_record(env={"GROK_HOME": tmp}, home=Path(tmp))
|
||
|
|
self.assertFalse(record["ready"])
|
||
|
|
self.assertIn("grok login", record["authHelpText"])
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|