#!/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()