broker/tests/test_creds_source.py

import contextlib
import hashlib
import io
import json
import os
import sys
import tempfile
import threading
import unittest
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)))

import coreci_broker

API_KEY = "sk-test-key"
KB_NAME = "creds-coreci-coo"
PROXMOX_SECRET = {"url": "https://pve.coreci.dev:8006", "token": "pve-tok"}
PROXMOX_UPDATED = "2026-10-01T12:00:00Z"


def _standard_fixtures():
    kbs = [{"id": "kb-1", "name": KB_NAME}]
    files = {"kb-1": [{"name": "coreci/proxmox", "id": "f1", "updated_at": PROXMOX_UPDATED}]}
    contents = {"f1": json.dumps(PROXMOX_SECRET)}
    return kbs, files, contents


class _StubOwuiHandler(BaseHTTPRequestHandler):
    kbs = None
    files = None
    contents = None
    failures = None
    hits = None

    def log_message(self, format, *args):
        pass

    def _send(self, status, payload):
        data = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        parsed = urllib.parse.urlsplit(self.path)
        path = parsed.path
        self.hits.append((path, self.headers.get("Authorization")))
        if path in self.failures:
            self._send(self.failures[path], {"error": self.failures[path]})
        elif path == "/api/v1/knowledge/":
            page = max(1, int((urllib.parse.parse_qs(parsed.query).get("page") or ["1"])[0]))
            self._send(200, {"items": self.kbs[(page - 1) * 30:page * 30], "total": len(self.kbs)})
        elif path.startswith("/api/v1/knowledge/"):
            kb_id = path[len("/api/v1/knowledge/"):]
            kb = next((entry for entry in self.kbs if entry["id"] == kb_id), None)
            if kb is None:
                self._send(404, {"error": 404})
            else:
                self._send(200, dict(kb, files=self.files.get(kb_id, [])))
        elif path.startswith("/api/v1/files/") and path.endswith("/data/content"):
            file_id = path[len("/api/v1/files/"):-len("/data/content")]
            if file_id in self.contents:
                self._send(200, {"content": self.contents[file_id]})
            else:
                self._send(404, {"error": 404})
        else:
            self._send(404, {"error": 404})


class StubOwuiServer:
    def __init__(self, kbs, files, contents):
        self.kbs, self.files, self.contents, self.failures = kbs, files, contents, {}
        self.hits = []
        handler = type("BoundStubHandler", (_StubOwuiHandler,), {
            "kbs": kbs, "files": files, "contents": contents,
            "failures": self.failures, "hits": self.hits,
        })
        self.server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()

    @property
    def base_url(self):
        return "http://127.0.0.1:%d" % self.server.server_address[1]

    def hits_for(self, path):
        return [hit for hit in self.hits if hit[0] == path]

    def close(self):
        self.server.shutdown()
        self.server.server_close()


class FileCredsSourceTests(unittest.TestCase):
    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.creds_path = os.path.join(tmp.name, "creds.json")
        self.document = {
            "coreci/proxmox": {"url": "https://pve.coreci.dev:8006", "token": "pve-tok"},
            "coreci/gitea": {"url": "https://git.coreci.dev", "token": "git-tok"},
        }
        self._write(self.document)
        self.source = coreci_broker.FileCredsSource(self.creds_path)

    def _write(self, document):
        with open(self.creds_path, "w", encoding="utf-8") as handle:
            json.dump(document, handle)

    def test_fetches_both_paths_with_response_shape(self):
        for path, secret in self.document.items():
            result = self.source.get_secret(path)
            self.assertEqual(path, result["path"])
            self.assertEqual(secret, result["data"])
            self.assertIsNone(result["updated_at"])
            self.assertRegex(result["version"], r"^[0-9a-f]{8}$")

    def test_version_is_short_sha256_of_serialized_data(self):
        result = self.source.get_secret("coreci/proxmox")
        serialized = json.dumps(result["data"], sort_keys=True, separators=(",", ":"))
        expected = hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:8]
        self.assertEqual(expected, result["version"])

    def test_equal_data_yields_equal_versions(self):
        first = self.source.get_secret("coreci/proxmox")
        second = self.source.get_secret("coreci/proxmox")
        self.assertEqual(first["version"], second["version"])

    def test_missing_path_is_not_found(self):
        with self.assertRaises(coreci_broker._SecretNotFound):
            self.source.get_secret("coreci/ghost")

    def test_missing_file_means_all_paths_not_found(self):
        os.remove(self.creds_path)
        for path in self.document:
            with self.assertRaises(coreci_broker._SecretNotFound):
                self.source.get_secret(path)

    def test_invalid_json_is_source_error(self):
        with open(self.creds_path, "w", encoding="utf-8") as handle:
            handle.write("{not-json")
        with self.assertRaises(coreci_broker._CredsSourceError):
            self.source.get_secret("coreci/proxmox")

    def test_non_object_document_is_source_error(self):
        self._write(["not", "a", "mapping"])
        with self.assertRaises(coreci_broker._CredsSourceError):
            self.source.get_secret("coreci/proxmox")

    def test_fresh_read_sees_out_of_band_edits(self):
        first = self.source.get_secret("coreci/proxmox")
        self.document["coreci/proxmox"]["token"] = "rotated"
        self._write(self.document)
        second = self.source.get_secret("coreci/proxmox")
        self.assertEqual("pve-tok", first["data"]["token"])
        self.assertEqual("rotated", second["data"]["token"])
        self.assertNotEqual(first["version"], second["version"])


class OwuiApiCredsSourceTests(unittest.TestCase):
    def setUp(self):
        kbs, files, contents = _standard_fixtures()
        self.stub = StubOwuiServer(kbs, files, contents)
        self.addCleanup(self.stub.close)
        self.source = coreci_broker.OwuiApiCredsSource(self.stub.base_url, API_KEY)

    def test_roundtrip_returns_response_shape(self):
        result = self.source.get_secret("coreci/proxmox")
        self.assertEqual("coreci/proxmox", result["path"])
        self.assertEqual(PROXMOX_SECRET, result["data"])
        self.assertEqual(PROXMOX_UPDATED, result["updated_at"])
        self.assertRegex(result["version"], r"^[0-9a-f]{8}$")

    def test_kb_and_file_map_cached_across_calls(self):
        self.source.get_secret("coreci/proxmox")
        self.source.get_secret("coreci/proxmox")
        self.assertEqual(1, len(self.stub.hits_for("/api/v1/knowledge/")))
        self.assertEqual(1, len(self.stub.hits_for("/api/v1/knowledge/kb-1")))
        self.assertEqual(2, len(self.stub.hits_for("/api/v1/files/f1/data/content")))

    def test_bearer_header_present_on_every_upstream_call(self):
        self.source.get_secret("coreci/proxmox")
        self.assertTrue(self.stub.hits)
        for _, auth in self.stub.hits:
            self.assertEqual("Bearer " + API_KEY, auth)

    def test_refresh_on_miss_picks_up_newly_added_files(self):
        self.source.get_secret("coreci/proxmox")
        self.stub.files["kb-1"].append(
            {"name": "coreci/gitea", "id": "f2", "updated_at": "2026-10-02T00:00:00Z"})
        self.stub.contents["f2"] = json.dumps({"user": "ceo"})
        result = self.source.get_secret("coreci/gitea")
        self.assertEqual({"user": "ceo"}, result["data"])
        self.assertEqual("2026-10-02T00:00:00Z", result["updated_at"])
        self.assertEqual(2, len(self.stub.hits_for("/api/v1/knowledge/")))

    def test_unknown_path_is_not_found(self):
        with self.assertRaises(coreci_broker._SecretNotFound):
            self.source.get_secret("coreci/ghost")

    def test_missing_kb_is_not_found(self):
        self.stub.kbs.clear()
        with self.assertRaises(coreci_broker._SecretNotFound):
            self.source.get_secret("coreci/proxmox")

    def test_kb_list_404_is_not_found(self):
        self.stub.failures["/api/v1/knowledge/"] = 404
        with self.assertRaises(coreci_broker._SecretNotFound):
            self.source.get_secret("coreci/proxmox")

    def test_kb_detail_404_is_not_found(self):
        self.stub.failures["/api/v1/knowledge/kb-1"] = 404
        with self.assertRaises(coreci_broker._SecretNotFound):
            self.source.get_secret("coreci/proxmox")

    def test_file_content_404_is_not_found(self):
        del self.stub.contents["f1"]
        with self.assertRaises(coreci_broker._SecretNotFound):
            self.source.get_secret("coreci/proxmox")

    def test_upstream_500_is_source_error_not_crash(self):
        self.stub.failures["/api/v1/files/f1/data/content"] = 500
        with self.assertRaises(coreci_broker._CredsSourceError):
            self.source.get_secret("coreci/proxmox")

    def test_kb_list_500_is_source_error(self):
        self.stub.failures["/api/v1/knowledge/"] = 500
        with self.assertRaises(coreci_broker._CredsSourceError):
            self.source.get_secret("coreci/proxmox")

    def test_unavailable_upstream_maps_to_distinct_source_error(self):
        source = coreci_broker.OwuiApiCredsSource("http://127.0.0.1:9", API_KEY)
        with self.assertRaises(coreci_broker._CredsSourceError) as caught:
            source.get_secret("coreci/proxmox")
        self.assertIn("unavailable", str(caught.exception))

    def test_invalid_content_json_is_source_error(self):
        self.stub.contents["f1"] = "not-json"
        with self.assertRaises(coreci_broker._CredsSourceError):
            self.source.get_secret("coreci/proxmox")

    def test_pagination_finds_kb_beyond_first_page(self):
        self.stub.kbs[:] = [{"id": "kb-%d" % i, "name": "kb-%d" % i} for i in range(30)]
        self.stub.kbs.append({"id": "kb-creds", "name": KB_NAME})
        self.stub.files["kb-creds"] = [
            {"name": "coreci/proxmox", "id": "f1", "updated_at": PROXMOX_UPDATED}]
        result = self.source.get_secret("coreci/proxmox")
        self.assertEqual(PROXMOX_SECRET, result["data"])
        self.assertEqual(2, len(self.stub.hits_for("/api/v1/knowledge/")))


class CredsSourceFactoryTests(unittest.TestCase):
    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.tmp = tmp.name

    def _boot(self, extra):
        env = {
            "BROKER_ENV_PATH": os.path.join(self.tmp, ".env"),
            "BROKER_DB_PATH": os.path.join(self.tmp, "broker.db"),
            "BROKER_AUDIT_LOG_PATH": os.path.join(self.tmp, "broker.audit.log"),
        }
        env.update(extra)
        with contextlib.redirect_stdout(io.StringIO()):
            broker = coreci_broker.boot(env)
        self.addCleanup(broker.audit.close)
        self.addCleanup(broker.store.close)
        return broker

    def test_default_source_is_file_with_default_path(self):
        broker = self._boot({})
        self.assertIsInstance(broker.creds_source, coreci_broker.FileCredsSource)
        self.assertEqual("data/creds.json", broker.creds_source.path)

    def test_explicit_file_source_with_custom_path(self):
        creds_path = os.path.join(self.tmp, "custom-creds.json")
        with open(creds_path, "w", encoding="utf-8") as handle:
            json.dump({"coreci/proxmox": PROXMOX_SECRET}, handle)
        broker = self._boot({
            "BROKER_CREDS_SOURCE": "file",
            "BROKER_CREDS_FILE": creds_path,
        })
        self.assertIsInstance(broker.creds_source, coreci_broker.FileCredsSource)
        self.assertEqual(creds_path, broker.creds_source.path)
        self.assertEqual(PROXMOX_SECRET, broker.creds_source.get_secret("coreci/proxmox")["data"])

    def test_owui_source_selected_by_config(self):
        kbs, files, contents = _standard_fixtures()
        stub = StubOwuiServer(kbs, files, contents)
        self.addCleanup(stub.close)
        broker = self._boot({
            "BROKER_CREDS_SOURCE": "owui",
            "BROKER_OWUI_BASE_URL": stub.base_url,
            "BROKER_OWUI_API_KEY": API_KEY,
        })
        source = broker.creds_source
        self.assertIsInstance(source, coreci_broker.OwuiApiCredsSource)
        self.assertEqual(stub.base_url, source.base_url)
        self.assertEqual(API_KEY, source.api_key)
        self.assertEqual(PROXMOX_SECRET, source.get_secret("coreci/proxmox")["data"])

    def test_unknown_source_value_raises_at_boot(self):
        with self.assertRaises(ValueError):
            self._boot({"BROKER_CREDS_SOURCE": "banana"})


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