codex-limit-monitor-waybar/tests/test_codex_waybar.py

import importlib.util
from importlib.machinery import SourceFileLoader
import io
import json
import os
import sys
import tempfile
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from unittest.mock import patch


def load_script(name, path):
    spec = importlib.util.spec_from_loader(name, SourceFileLoader(name, str(path)))
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


ROOT = Path(__file__).resolve().parents[1]
module = load_script("codex_waybar", ROOT / "bin" / "codex-waybar")
configure = load_script("codex_waybar_configure", ROOT / "bin" / "codex-waybar-configure")


def rate_result(five_used=22, weekly_used=18):
    return {
        "rateLimits": {
            "limitId": "codex",
            "planType": "plus",
            "primary": {
                "usedPercent": five_used,
                "windowDurationMins": 300,
                "resetsAt": 10_000,
            },
            "secondary": {
                "usedPercent": weekly_used,
                "windowDurationMins": 10_080,
                "resetsAt": 200_000,
            },
        }
    }


class UsageModuleTests(unittest.TestCase):
    def test_parse_json_lines_ignores_notifications(self):
        output = "\n".join(
            [
                '{"id":1,"result":{}}',
                '{"method":"notification","params":{}}',
                json.dumps({"id": 2, "result": rate_result()}),
            ]
        )
        self.assertEqual(module.parse_json_lines(output), rate_result())

    def test_selects_windows_by_duration_not_position(self):
        result = rate_result()
        snapshot = result["rateLimits"]
        snapshot["primary"], snapshot["secondary"] = snapshot["secondary"], snapshot["primary"]
        values, percentage = module.format_values(snapshot, now=1_000)
        self.assertEqual(values["five_hour_used"], 22)
        self.assertEqual(values["weekly_used"], 18)
        self.assertEqual(percentage, 22)

    def test_countdown_formats_and_clamps(self):
        self.assertEqual(module.human_countdown(1_000, 1_000), "0m")
        self.assertEqual(module.human_countdown(4_601, 1_000), "1h 01m")
        self.assertEqual(module.human_countdown(91_000, 1_000), "1d 1h")
        self.assertEqual(module.human_countdown(999, 1_000), "0m")

    def test_output_sets_classes_and_custom_format(self):
        output = module.output_json(
            rate_result(five_used=72),
            "{five_hour_used}/{weekly_used} {five_hour_reset_in}",
            "{plan}",
            warning=70,
            critical=90,
            now=1_000,
        )
        self.assertEqual(output["text"], "72/18 2h 30m")
        self.assertEqual(output["tooltip"], "plus")
        self.assertEqual(output["class"], ["warning"])
        self.assertEqual(output["percentage"], 72)

    def test_missing_window_is_rejected(self):
        result = rate_result()
        result["rateLimits"]["secondary"] = None
        with self.assertRaises(module.UsageError):
            module.output_json(result, module.DEFAULT_FORMAT, module.DEFAULT_TOOLTIP, 70, 90)

    def test_unknown_template_field_is_rejected(self):
        with self.assertRaises(module.UsageError):
            module.validate_template("{not_a_field}")

    def test_main_uses_stale_cache_after_fetch_failure(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "codex-waybar" / "rate-limits.json"
            module.save_cache(path, rate_result())
            stdout = io.StringIO()
            with (
                patch.dict(os.environ, {"XDG_CACHE_HOME": tmp}),
                patch.object(module, "fetch_rate_limits", side_effect=module.UsageError("offline")),
                redirect_stdout(stdout),
            ):
                exit_code = module.main(["--format", "{five_hour_used}%"])
            output = json.loads(stdout.getvalue())
            self.assertEqual(exit_code, 0)
            self.assertEqual(output["text"], "22%")
            self.assertIn("stale", output["class"])

    def test_main_returns_unavailable_without_cache(self):
        with tempfile.TemporaryDirectory() as tmp:
            stdout = io.StringIO()
            with (
                patch.dict(os.environ, {"XDG_CACHE_HOME": tmp}),
                patch.object(module, "fetch_rate_limits", side_effect=module.UsageError("offline")),
                redirect_stdout(stdout),
            ):
                module.main([])
            output = json.loads(stdout.getvalue())
            self.assertEqual(output["class"], ["unavailable"])
            self.assertIn("offline", output["tooltip"])


class ConfigureTests(unittest.TestCase):
    CONFIG = """{
  // Existing comment
  "modules-left": ["clock"],
  "modules-center": [],
  "modules-right": [
    "cpu"
  ],
  "cpu": {"interval": 5}
}
"""

    def test_configure_preserves_comments_and_is_idempotent(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Path(tmp) / "config.jsonc"
            style = Path(tmp) / "style.css"
            config.write_text(self.CONFIG)
            style.write_text("* { color: white; }\n")

            configure.configure(config, style, "right")
            first_config = config.read_text()
            first_style = style.read_text()
            configure.configure(config, style, "right")

            self.assertIn("// Existing comment", first_config)
            self.assertEqual(first_config.count('"custom/codex-usage"'), 2)
            self.assertEqual(config.read_text(), first_config)
            self.assertEqual(style.read_text(), first_style)
            self.assertEqual(first_style.count("#custom-codex-usage"), 5)

    def test_ambiguous_section_aborts_before_write(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Path(tmp) / "config.jsonc"
            style = Path(tmp) / "style.css"
            original = '{"modules-right": [], "modules-right": []}\n'
            config.write_text(original)
            style.write_text("")
            with self.assertRaises(configure.ConfigureError):
                configure.configure(config, style, "right")
            self.assertEqual(config.read_text(), original)


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