user-engine/tests/test_custody_transport.py

import io
import json
from pathlib import Path
import tempfile
import unittest
from datetime import datetime, timezone
from urllib.error import HTTPError

from user_engine.custody_runtime import OpenBaoCustodyReader, CustodyReadUnavailable


class Response(io.BytesIO):
    def __init__(self, value, *, headers=None):
        super().__init__(json.dumps(value).encode())
        self.headers = headers or {}
        self.status = 200


class Opener:
    def __init__(self, responses):
        self.responses, self.requests = list(responses), []
    def open(self, request, timeout):
        self.requests.append((request, timeout))
        response = self.responses.pop(0)
        if isinstance(response, Exception):
            raise response
        return response


class CustodyTransportTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.jwt = Path(self.temp.name) / "token"
        self.jwt.write_text("fixture-jwt\n")
        self.now = datetime.now(timezone.utc)
        self.document = {"revision": 4}

    def reader(self, opener, **options):
        return OpenBaoCustodyReader(base_url="https://bao.example", auth_mount="kubernetes",
            role="user-engine-root-custody-reader", jwt_file=str(self.jwt), source_id="trusted",
            opener=opener, clock=lambda: self.now, **options)

    def login(self):
        return Response({"auth": {"client_token": "fixture-reader-token"}})

    def test_dedicated_login_fresh_exact_read_and_revoke(self):
        opener = Opener([self.login(), Response({"data": {"data": {"document": self.document}}}), Response({})])
        result = self.reader(opener).read()
        self.assertEqual(result.document, self.document)
        self.assertEqual(result.read_at, self.now)
        login, read, revoke = [entry[0] for entry in opener.requests]
        self.assertEqual(json.loads(login.data), {"role": "user-engine-root-custody-reader", "jwt": "fixture-jwt"})
        self.assertEqual(read.full_url, "https://bao.example/v1/platform/data/custody/net-kingdom/platform-root")
        self.assertEqual(read.get_header("X-vault-token"), "fixture-reader-token")
        self.assertTrue(revoke.full_url.endswith("auth/token/revoke-self"))
        self.assertTrue(all(0 < timeout <= 3 for _, timeout in opener.requests))

    def test_owner_denial_errors_are_suppressed(self):
        opener = Opener([HTTPError("https://secret-value.invalid", 403, "sensitive", {}, None)])
        with self.assertRaises(CustodyReadUnavailable) as error:
            self.reader(opener).read()
        self.assertEqual(str(error.exception), "custody read unavailable")
        self.assertIsNone(error.exception.__cause__)

    def test_cache_hit_malformed_and_oversized_responses_deny(self):
        for response in (Response({"data": {}}, headers={"Age": "1"}), Response({"wrong": "envelope"}),
                         Response({"data": {"data": {"document": {"padding": "x" * 2000}}}})):
            opener = Opener([self.login(), response, Response({})])
            with self.assertRaises(CustodyReadUnavailable):
                self.reader(opener, max_bytes=1024).read()

    def test_token_rotation_is_read_from_file_each_login(self):
        opener = Opener([self.login(), Response({"data": {"data": {"document": self.document}}}), Response({}),
                         self.login(), Response({"data": {"data": {"document": self.document}}}), Response({})])
        reader = self.reader(opener)
        reader.read()
        self.jwt.write_text("rotated-fixture-jwt")
        reader.read()
        self.assertEqual(json.loads(opener.requests[3][0].data)["jwt"], "rotated-fixture-jwt")

    def test_plaintext_untrusted_origin_and_path_are_configuration_errors(self):
        for origin in ("http://bao.example", "https://user:pass@bao.example", "https://bao.example/path"):
            with self.assertRaises(ValueError):
                OpenBaoCustodyReader(base_url=origin, auth_mount="kubernetes", role="reader", jwt_file="fixture", source_id="source")
        with self.assertRaises(ValueError):
            self.reader(Opener([]), data_path="platform/data/other")


class CustodyNestedIntegrationTests(unittest.TestCase):
    def test_nested_owner_response_to_serialized_account_authority(self):
        from datetime import timedelta
        from dataclasses import replace
        from user_engine.adapters.local import InMemoryUserEngineStore
        from user_engine.custody import CustodyObservationAdapter
        from user_engine.custody_canon import policy
        from user_engine.service import UserEngineService
        from user_engine.testing.fixtures import FixtureIdentityClaimsAdapter, StaticAuthorizationCheckPort, human_actor_claims
        from user_engine.web import _jsonable
        from test_custody_adapter import Baselines, CANON
        if not CANON.exists():
            self.skipTest("optional sibling fixtures absent")
        document = json.loads((CANON / "fixtures/valid/grant.json").read_text())
        now = datetime(2026, 10, 3, 12, tzinfo=timezone.utc)
        claims = human_actor_claims()
        claims.update(iss=document["entries"][0]["subject"]["issuer"], sub=document["entries"][0]["subject"]["subject"], tenant="platform")
        with tempfile.TemporaryDirectory() as directory:
            jwt_file = Path(directory) / "token"
            jwt_file.write_text("fixture-openbao-audience-jwt")
            opener = Opener([Response({"auth": {"client_token": "fixture-token"}}),
                Response({"data": {"data": {"document": document}}}), Response({})])
            source = "https://bao.example/v1/platform/data/custody/net-kingdom/platform-root#auth=kubernetes/user-engine-root-custody-reader"
            reader = OpenBaoCustodyReader(base_url="https://bao.example", auth_mount="kubernetes",
                role="user-engine-root-custody-reader", jwt_file=str(jwt_file), source_id=source,
                opener=opener, clock=lambda: now)
            adapter = CustodyObservationAdapter(reader=reader, baselines=Baselines(), policy=policy,
                source_id=source, record_id=document["record_id"], allow_initial_baseline=True,
                clock=lambda: now + timedelta(milliseconds=1))
            service = UserEngineService(store=InMemoryUserEngineStore(), identity_adapter=FixtureIdentityClaimsAdapter(),
                authorization=StaticAuthorizationCheckPort(), custody=adapter)
            session = service.me(claims)
            result = service.lookup_account_authority(replace(session.actor, subject="fixture-hub-reader"),
                issuer=claims["iss"], subject=claims["sub"], tenant="platform")
            envelope = _jsonable(result)
            observation = envelope["root_grant"]
            self.assertEqual(observation["source"], {"record_id": document["record_id"], "revision": 1})
            self.assertEqual(observation["subject"], document["entries"][0]["subject"])
            self.assertEqual(observation["observed_at"], "2026-10-03T12:00:00Z")
            self.assertEqual(observation["expires_at"], "2026-10-03T12:15:00Z")
            self.assertEqual(opener.requests[1][0].full_url, "https://bao.example/v1/platform/data/custody/net-kingdom/platform-root")