DocForge2/tests/test_visualization.py

from __future__ import annotations

import json
import shutil
import tempfile
import unittest
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

from docforge.errors import DocForgeError
from docforge.index import ProjectIndex
from docforge.project import Project
from docforge.visualization import (
    VISUALIZATION_TEMPLATE,
    VisualizationIndexSnapshot,
    VisualizationRunner,
)

ROOT = Path(__file__).resolve().parents[1]
FIXTURES = ROOT / "tests" / "fixtures"


class VisualizationTests(unittest.TestCase):
    def copy_fixture(self, name: str, destination: Path) -> Path:
        root = destination / name
        shutil.copytree(FIXTURES / name, root)
        return root

    def test_overview_and_neighborhood_are_deterministic_and_bounded(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            index = ProjectIndex(Project.open(root))
            index.build()
            snapshot = VisualizationIndexSnapshot(index, index.check())

            overview = snapshot.overview()
            first = snapshot.node("guide.workflow", depth=2, limit=2)
            second = snapshot.node("guide.workflow", depth=2, limit=2)

            self.assertEqual(3, overview["node_count"])
            self.assertEqual(2, overview["edge_count"])
            self.assertEqual(
                [
                    {"value": "guide", "count": 2},
                    {"value": "proof", "count": 1},
                ],
                overview["families"],
            )
            self.assertEqual(first, second)
            self.assertEqual("guide.workflow", first["root"])
            self.assertLessEqual(len(first["edges"]), 2)
            self.assertIn(
                "guide.workflow",
                {node["node_id"] for node in first["nodes"]},
            )

            with self.assertRaisesRegex(DocForgeError, "safety boundary"):
                snapshot.node("guide.workflow", depth=1, limit=401)

    def test_runner_serves_only_token_bound_read_only_graph_endpoints(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            index = ProjectIndex(Project.open(root))
            index.build()
            runner = VisualizationRunner(index)
            try:
                first = runner.start(node_id="guide.workflow", depth=2)
                second = runner.start(query="canonical nodes")
                first_url = urllib.parse.urlparse(str(first["url"]))
                second_url = urllib.parse.urlparse(str(second["url"]))
                base = f"{first_url.scheme}://{first_url.netloc}{first_url.path}"

                self.assertEqual(VISUALIZATION_TEMPLATE, first["template"])
                self.assertTrue(first["read_only"])
                self.assertEqual(first_url.netloc, second_url.netloc)
                self.assertEqual(first_url.path, second_url.path)

                with urllib.request.urlopen(base, timeout=2) as response:
                    html = response.read().decode("utf-8")
                    headers = response.headers
                self.assertIn("DocForge graph", html)
                self.assertIn('id="zoom-in"', html)
                self.assertIn('id="zoom-out"', html)
                self.assertIn('id="reset-view"', html)
                self.assertIn('addEventListener("wheel"', html)
                self.assertIn('addEventListener("pointermove"', html)
                self.assertIn("left-drag to pan", html)
                self.assertIn("default-src 'none'", headers["Content-Security-Policy"])
                self.assertEqual("no-store", headers["Cache-Control"])
                self.assertEqual("DENY", headers["X-Frame-Options"])

                with urllib.request.urlopen(f"{base}api/overview", timeout=2) as response:
                    overview = json.load(response)
                self.assertEqual("alpha-docs", overview["project_id"])
                self.assertEqual(3, overview["node_count"])
                self.assertEqual(20, overview["max_results"])

                search_query = urllib.parse.urlencode(
                    {"q": "canonical nodes", "family": "", "limit": overview["max_results"]}
                )
                with urllib.request.urlopen(
                    f"{base}api/search?{search_query}", timeout=2
                ) as response:
                    search = json.load(response)
                self.assertEqual("alpha-docs", search["project_id"])
                self.assertGreaterEqual(search["count"], 1)

                node_query = urllib.parse.urlencode(
                    {"id": "guide.workflow", "depth": 1, "limit": 20}
                )
                with urllib.request.urlopen(f"{base}api/node?{node_query}", timeout=2) as response:
                    node = json.load(response)
                self.assertEqual("guide.workflow", node["node"]["node_id"])
                self.assertEqual("guide.workflow", node["root"])

                wrong_token = f"{first_url.scheme}://{first_url.netloc}/wrong-token/api/overview"
                with self.assertRaises(urllib.error.HTTPError) as missing:
                    urllib.request.urlopen(wrong_token, timeout=2)
                self.assertEqual(404, missing.exception.code)
                missing.exception.close()

                request = urllib.request.Request(
                    f"{base}api/overview",
                    data=b"{}",
                    method="POST",
                )
                with self.assertRaises(urllib.error.HTTPError) as rejected:
                    urllib.request.urlopen(request, timeout=2)
                self.assertEqual(405, rejected.exception.code)
                try:
                    self.assertEqual(
                        "method_not_allowed",
                        json.loads(rejected.exception.read())["error"]["code"],
                    )
                finally:
                    rejected.exception.close()
            finally:
                runner.stop()

    def test_runner_rejects_ambiguous_targets_and_changed_index_snapshot(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            index = ProjectIndex(Project.open(root))
            index.build()
            runner = VisualizationRunner(index)
            try:
                with self.assertRaisesRegex(DocForgeError, "either one exact"):
                    runner.start(node_id="guide.workflow", query="workflow")
                result = runner.start()
                parsed = urllib.parse.urlparse(str(result["url"]))
                with index.path.open("ab") as handle:
                    handle.write(b"\n")
                endpoint = f"{parsed.scheme}://{parsed.netloc}{parsed.path}api/overview"
                with self.assertRaises(urllib.error.HTTPError) as stale:
                    urllib.request.urlopen(endpoint, timeout=2)
                self.assertEqual(409, stale.exception.code)
                try:
                    self.assertEqual(
                        "visualization_stale",
                        json.loads(stale.exception.read())["error"]["code"],
                    )
                finally:
                    stale.exception.close()
            finally:
                runner.stop()

    def test_two_visualizations_remain_project_bound(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            parent = Path(directory)
            alpha_root = self.copy_fixture("alpha", parent / "alpha")
            beta_root = self.copy_fixture("beta", parent / "beta")
            alpha_index = ProjectIndex(Project.open(alpha_root))
            beta_index = ProjectIndex(Project.open(beta_root))
            alpha_index.build()
            beta_index.build()
            alpha = VisualizationRunner(alpha_index)
            beta = VisualizationRunner(beta_index)
            try:
                alpha_url = urllib.parse.urlparse(str(alpha.start()["url"]))
                beta_url = urllib.parse.urlparse(str(beta.start()["url"]))
                self.assertNotEqual(alpha_url.netloc, beta_url.netloc)
                for parsed, expected in (
                    (alpha_url, "alpha-docs"),
                    (beta_url, "beta-notes"),
                ):
                    endpoint = f"{parsed.scheme}://{parsed.netloc}{parsed.path}api/overview"
                    with urllib.request.urlopen(endpoint, timeout=2) as response:
                        self.assertEqual(expected, json.load(response)["project_id"])
            finally:
                alpha.stop()
                beta.stop()


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