test_opencode_usage_monitor.py

import contextlib
import io
import json
import os
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch
from datetime import datetime, timedelta, timezone

import opencode_usage_monitor as m


def fixture(percent=9):
    return {'usage': {w: {'percent': percent, 'status': 'ok',
                          'resetsAt': '2030-01-01T00:00:00Z'} for w in m.WINDOWS}}


class WatchdogTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.root = Path(self.tmp.name)
        env = patch.dict(os.environ, {'HOME': str(self.root), 'HERMES_HOME': str(self.root),
                                     'OPENCODE_GO_API_KEY': 'secret-alpha'}, clear=True)
        env.start()
        self.addCleanup(env.stop)

    def config(self, data):
        p = self.root / 'keys.json'
        p.write_text(json.dumps(data))
        return str(p)

    def run_main(self, args):
        out, err = io.StringIO(), io.StringIO()
        with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
            rc = m.main(args)
        return rc, out.getvalue(), err.getvalue()

    def test_duplicate_keys(self):
        p = self.config([{'name': 'A', 'key': 'secret-alpha'},
                         {'name': 'B', 'key': 'secret-alpha'},
                         {'name': 'C', 'key': 'secret-beta'}])
        self.assertEqual(len(m.load_keys(p)), 2)

    def test_bad_config_no_fallback(self):
        for data in ([], {}, [1], [{'name': 'A'}], [{'name': '\n', 'key': 'abc'}],
                     [{'name': 'secret-alpha', 'key': 'secret-alpha'}]):
            with self.subTest(data=data):
                with self.assertRaises(ValueError):
                    m.load_keys(self.config(data))
        p = self.root / 'bad.json'
        p.write_text('{')
        with self.assertRaises(ValueError):
            m.load_keys(str(p))
        with self.assertRaises(ValueError):
            m.load_keys(str(self.root / 'missing'))

    def test_single_env_discovery(self):
        self.assertEqual(m.get_key(), 'secret-alpha')
        del os.environ['OPENCODE_GO_API_KEY']
        os.environ['OPENCODE_ZEN_API_KEY'] = 'zen-key'
        self.assertEqual(m.get_key(), 'zen-key')
        del os.environ['OPENCODE_ZEN_API_KEY']
        (self.root / '.env').write_text('export OPENCODE_GO_API_KEY="file-key"\n')
        self.assertEqual(m.get_key(), 'file-key')
        self.assertEqual(len(m.load_keys()), 1)

    def test_state_isolation(self):
        a, ah = m.paths_for_key('secret-alpha')
        b, bh = m.paths_for_key('secret-beta')
        self.assertNotEqual(a, b)
        self.assertNotEqual(ah, bh)
        self.assertIn(m.key_id('secret-alpha'), a)
        self.assertNotIn('secret-alpha', a)
        self.assertTrue(m.tick(fixture(), state_path=a, history_path=ah))
        self.assertIsNone(m.tick(fixture(), state_path=a, history_path=ah))
        self.assertTrue(m.tick(fixture(), state_path=b, history_path=bh))

    def test_once_no_writes(self):
        with patch.object(m, 'fetch', return_value=fixture()):
            rc, out, _ = self.run_main(['--once'])
        self.assertEqual(rc, 0)
        self.assertIn('9%', out)
        self.assertEqual(list(self.root.iterdir()), [])

    def test_partial_failure_json_redaction(self):
        p = self.config([{'name': 'A', 'key': 'secret-alpha'}, {'name': 'B', 'key': 'secret-beta'}])
        def fetch(key):
            if key == 'secret-beta':
                raise RuntimeError('Authorization: secret-beta body secret-alpha')
            data = fixture()
            data['usage']['rolling']['status'] = 'secret-alpha'
            data['usage']['channel'] = 'secret-beta'
            return data
        with patch.object(m, 'fetch', side_effect=fetch):
            rc, out, err = self.run_main(['--keys-file', p, '--once', '--json'])
        self.assertNotEqual(rc, 0)
        result = json.loads(out)
        self.assertEqual(result['count'], 2)
        self.assertEqual(result['succeeded'], 1)
        self.assertEqual(result['failed'], 1)
        self.assertEqual([r['success'] for r in result['rows']], [True, False])
        self.assertTrue(all(r['timestamp'] for r in result['rows']))
        self.assertNotIn('secret-', out + err)

    def test_percent_schema_and_units(self):
        for percent in (0, 1.5, 125):
            self.assertEqual(m.validate_usage(fixture(percent))['usage']['rolling']['percent'], percent)
        for percent in (-1, True, '4', float('nan'), float('inf')):
            with self.assertRaises(ValueError):
                m.validate_usage(fixture(percent))
        data = fixture()
        for field, value in [('status', 1), ('resetsAt', '2030-01-01'), ('resetsAt', 'secret-alpha')]:
            bad = fixture()
            bad['usage']['rolling'][field] = value
            with self.assertRaises(ValueError):
                m.validate_usage(bad)
        out = m.build_report(data, forced=True)
        self.assertNotIn('$', out)
        self.assertNotIn('token', out.lower())
        self.assertIn('9%', out)

    def test_thresholds_reset_daily(self):
        s, h = m.paths_for_key('secret-alpha')
        now = datetime(2026, 1, 1, tzinfo=timezone.utc)
        def tick(data, t=now):
            return m.tick(data, now=t, state_path=s, history_path=h)
        self.assertTrue(tick(fixture()))
        self.assertIsNone(tick(fixture()))
        self.assertTrue(tick(fixture(11)))
        self.assertTrue(tick(fixture(70)))
        self.assertTrue(tick(fixture(85)))
        self.assertTrue(tick(fixture(9)))
        data = fixture()
        data['usage']['rolling']['resetsAt'] = '2030-01-02T00:00:00Z'
        # Idle accounts may slide resetsAt every poll: that alone is not an alert.
        self.assertIsNone(tick(data))
        self.assertIsNone(tick(data))
        self.assertTrue(tick(data, now + timedelta(days=1)))
        self.assertIsNone(tick(data, now + timedelta(days=1)))

    def test_failure_cooldown(self):
        s, _ = m.paths_for_key('secret-alpha')
        now = datetime(2026, 1, 1, tzinfo=timezone.utc)
        first = m.tick_failure('secret-alpha', now=now, state_path=s)
        self.assertTrue(first)
        self.assertNotIn('secret-alpha', first)
        self.assertIsNone(m.tick_failure('secret-alpha', now=now, state_path=s))
        self.assertTrue(m.tick_failure('x', now=now + timedelta(hours=2), state_path=s))

    def test_redirect_and_response_limit(self):
        with self.assertRaises(ValueError):
            m.NoRedirect().redirect_request(None, None, 302, '', {}, 'https://other.invalid')
        from unittest.mock import MagicMock
        opener = MagicMock()
        opener.open.return_value.__enter__.return_value.read.return_value = b' ' * (1024 * 1024 + 1)
        with patch.object(m.urllib.request, 'build_opener', return_value=opener):
            with self.assertRaises(ValueError):
                m.fetch('secret-alpha')
        request = opener.open.call_args.args[0]
        self.assertEqual(request.get_header('Authorization'), 'Bearer secret-alpha')
        self.assertIn('Mozilla', request.get_header('User-agent'))

    def test_cross_key_label_rejected(self):
        p = self.config([{'name': 'secret-beta', 'key': 'secret-alpha'},
                         {'name': 'B', 'key': 'secret-beta'}])
        with patch.object(m, 'fetch') as fetch:
            rc, out, err = self.run_main(['--keys-file', p, '--json'])
        self.assertEqual(rc, 1)
        fetch.assert_not_called()
        self.assertNotIn('secret-', out + err)
        self.assertEqual(json.loads(out)['count'], 0)

    def test_monitor_failure_exit_and_cooldown(self):
        with patch.object(m, 'fetch', side_effect=RuntimeError('secret-alpha')):
            rc, out, err = self.run_main([])
            self.assertEqual(rc, 2)
            self.assertTrue(out)
            self.assertNotIn('secret-alpha', out + err)
            rc, out, _ = self.run_main([])
            self.assertEqual(rc, 2)
            self.assertEqual(out, '')
            rc, out, _ = self.run_main(['--json'])
            self.assertEqual(rc, 2)
            self.assertEqual(json.loads(out)['failed'], 1)

    def test_json_silent_success_and_cumulative_delta(self):
        with patch.object(m, 'fetch', return_value=fixture()):
            self.run_main([])
            rc, out, _ = self.run_main(['--json'])
        self.assertEqual(rc, 0)
        self.assertEqual(json.loads(out)['count'], 1)
        s, h = m.paths_for_key('secret-alpha')
        self.assertIsNone(m.tick(fixture(10), state_path=s, history_path=h))
        self.assertTrue(m.tick(fixture(11), state_path=s, history_path=h))

    def test_worker_concurrency_is_bounded(self):
        import threading
        import time
        p = self.config([{'name': str(i), 'key': 'synthetic-key-' + str(i)} for i in range(6)])
        active = peak = 0
        lock = threading.Lock()
        def fetch(key):
            nonlocal active, peak
            with lock:
                active += 1
                peak = max(peak, active)
            time.sleep(0.02)
            with lock:
                active -= 1
            return fixture()
        with patch.object(m, 'fetch', side_effect=fetch):
            rc, out, _ = self.run_main(['--keys-file', p, '--once', '--json', '--workers', '2'])
        self.assertEqual(rc, 0)
        self.assertEqual(peak, 2)
        self.assertEqual(json.loads(out)['count'], 6)

    def test_workers_bounds(self):
        self.assertEqual(m.parse_args([]).workers, 4)
        for n in ('0', '17'):
            with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
                m.parse_args(['--workers', n])


if __name__ == '__main__':
    unittest.main()