cc-os/plugins/os-context/tests/hook_test.py

"""Tests for hooks/inject.py. Run: python3 tests/hook_test.py"""

import json
import sys
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

PLUGIN_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PLUGIN_ROOT / "hooks"))

import inject  # noqa: E402


def write_prompts(directory, files):
    """files: dict of filename -> content, written under directory."""
    directory.mkdir(parents=True, exist_ok=True)
    for name, content in files.items():
        (directory / name).write_text(content)


class ComposeOrderTest(unittest.TestCase):
    def test_concatenates_in_filename_sort_order(self):
        with TemporaryDirectory() as tmp:
            d = Path(tmp)
            write_prompts(d, {"20-second.md": "second\n", "10-first.md": "first\n"})
            paths = inject.prompt_files(d)
            self.assertEqual(["10-first.md", "20-second.md"], [Path(p).name for p in paths])
            self.assertEqual("first\nsecond\n", inject.compose(paths))

    def test_no_separator_added_between_files(self):
        with TemporaryDirectory() as tmp:
            d = Path(tmp)
            write_prompts(d, {"10-a.md": "a", "20-b.md": "b"})
            paths = inject.prompt_files(d)
            self.assertEqual("ab", inject.compose(paths))


class ByteIdentityTest(unittest.TestCase):
    """With only the real 10-orchestration.md present, the composer's output
    must be byte-identical to what the old single-file inject.py emitted
    directly from ORCHESTRATION.md's contents — no added separators or
    trailing newlines."""

    def test_single_file_matches_its_raw_contents_exactly(self):
        real_file = PLUGIN_ROOT / "prompts" / "session-start" / "10-orchestration.md"
        with TemporaryDirectory() as tmp:
            d = Path(tmp)
            d.mkdir(exist_ok=True)
            shadow = d / "10-orchestration.md"
            shadow.write_bytes(real_file.read_bytes())
            paths = inject.prompt_files(d)
            composed = inject.compose(paths)
            self.assertEqual(real_file.read_text(), composed)

    def test_real_prompts_dir_composes_all_files_byte_identical_concatenation(self):
        paths = inject.prompt_files()
        self.assertGreaterEqual(len(paths), 1)
        expected = "".join(Path(p).read_text() for p in paths)
        self.assertEqual(expected, inject.compose(paths))


class BudgetTest(unittest.TestCase):
    def make_dir_with_lines(self, tmp, n_lines):
        d = Path(tmp)
        d.mkdir(exist_ok=True)
        content = "\n".join(f"line {i}" for i in range(n_lines)) + "\n"
        (d / "10-big.md").write_text(content)
        return d

    def test_under_warn_budget_injects_silently(self):
        with TemporaryDirectory() as tmp:
            d = self.make_dir_with_lines(tmp, inject.WARN_LINE_BUDGET - 1)
            paths = inject.prompt_files(d)
            context = inject.compose(paths)
            self.assertLessEqual(len(context.splitlines()), inject.WARN_LINE_BUDGET)

    def test_over_warn_under_refuse_still_injects(self):
        with TemporaryDirectory() as tmp:
            d = self.make_dir_with_lines(
                tmp, inject.WARN_LINE_BUDGET + 1
            )
            paths = inject.prompt_files(d)
            context = inject.compose(paths)
            line_count = len(context.splitlines())
            self.assertGreater(line_count, inject.WARN_LINE_BUDGET)
            self.assertLessEqual(line_count, inject.REFUSE_LINE_BUDGET)
            # Injection still happens: compose() returns the full content.
            self.assertTrue(context)

    def test_over_refuse_budget_line_count_detected(self):
        with TemporaryDirectory() as tmp:
            d = self.make_dir_with_lines(tmp, inject.REFUSE_LINE_BUDGET + 1)
            paths = inject.prompt_files(d)
            context = inject.compose(paths)
            self.assertGreater(len(context.splitlines()), inject.REFUSE_LINE_BUDGET)

    def test_main_warns_on_stderr_between_warn_and_refuse(self, ):
        import io
        from unittest import mock

        with TemporaryDirectory() as tmp:
            d = self.make_dir_with_lines(tmp, inject.WARN_LINE_BUDGET + 1)
            stderr = io.StringIO()
            stdout = io.StringIO()
            with mock.patch.object(inject, "SESSION_START_DIR", str(d)), \
                 mock.patch("sys.stderr", stderr), \
                 mock.patch("sys.stdout", stdout):
                inject.main()
            self.assertIn("injecting anyway", stderr.getvalue())
            payload = json.loads(stdout.getvalue())
            self.assertIn("additionalContext", payload["hookSpecificOutput"])

    def test_main_refuses_and_emits_nothing_over_refuse_budget(self):
        import io
        from unittest import mock

        with TemporaryDirectory() as tmp:
            d = self.make_dir_with_lines(tmp, inject.REFUSE_LINE_BUDGET + 1)
            stderr = io.StringIO()
            stdout = io.StringIO()
            with mock.patch.object(inject, "SESSION_START_DIR", str(d)), \
                 mock.patch("sys.stderr", stderr), \
                 mock.patch("sys.stdout", stdout):
                inject.main()
            self.assertIn("refusing to inject", stderr.getvalue())
            self.assertEqual("", stdout.getvalue())


class NoPromptsTest(unittest.TestCase):
    def test_main_emits_nothing_when_dir_empty(self):
        import io
        from unittest import mock

        with TemporaryDirectory() as tmp:
            d = Path(tmp) / "empty"
            d.mkdir()
            stdout = io.StringIO()
            with mock.patch.object(inject, "SESSION_START_DIR", str(d)), \
                 mock.patch("sys.stdout", stdout):
                inject.main()
            self.assertEqual("", stdout.getvalue())


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