teahub/mcp/tests/test_domains.py

"""GET /domains against a fake Cloudflare. From mcp/: python -m unittest discover -s tests"""
from __future__ import annotations

import json
import os
import sys
import unittest
from unittest import mock

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

import httpx  # noqa: E402
from starlette.applications import Starlette  # noqa: E402
from starlette.routing import Route  # noqa: E402

from tea_mcp import domains  # noqa: E402

ORIGIN = "https://teahub.test"


def offer(name: str, registrable: bool = True, cost: str = "10.46") -> dict:
    if not registrable:
        return {"name": name, "registrable": False, "tier": "standard", "reason": "domain_unavailable"}
    return {"name": name, "registrable": True, "tier": "standard",
            "pricing": {"currency": "USD", "registration_cost": cost, "renewal_cost": cost}}


class FakeCloudflare:
    def __init__(self):
        self.calls: list[httpx.Request] = []
        self.fail = 0

    def __call__(self, request: httpx.Request) -> httpx.Response:
        self.calls.append(request)
        if request.headers.get("authorization") != "Bearer cf-token" or "/accounts/acct/" not in request.url.path:
            return httpx.Response(403, json={"success": False, "errors": [{"code": 10000}]})
        if self.fail:
            return httpx.Response(self.fail, json={"success": False, "errors": [{"code": 1000}]})
        if request.url.path.endswith("/registrar/domain-search"):
            return httpx.Response(200, json={"success": True, "result": {"domains": [
                offer("teahub.net"), offer("teahubusa.com"), offer("teahub.uk", cost="5.30")]}})
        names = json.loads(request.content)["domains"]
        if names == ["a.b"]:
            return httpx.Response(400, json={"success": False, "errors": [{"code": 1008}]})
        return httpx.Response(200, json={"success": True, "result": {"domains": [
            offer(name, registrable=not name.endswith(".com")) for name in names]}})


class Domains(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        env = mock.patch.dict(os.environ, {"GITEA_URL": ORIGIN, "CLOUDFLARE_API_TOKEN": "cf-token",
                                           "CLOUDFLARE_ACCOUNT_ID": "acct"})
        env.start()
        self.addCleanup(env.stop)
        self.cf = FakeCloudflare()
        domains.transport = httpx.MockTransport(self.cf)
        self.addCleanup(setattr, domains, "transport", None)
        for kept in (domains._kept, domains._by_client, domains._by_all):
            kept.clear()
        app = Starlette(routes=[Route("/domains", domains.domains_http, methods=["GET", "OPTIONS"])])
        self.web = httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="https://mcp.test")
        self.addAsyncCleanup(self.web.aclose)

    async def ask(self, q: str, origin: str = ORIGIN, ip: str = "203.0.113.7", **params) -> httpx.Response:
        return await self.web.get("/domains", params={"q": q, **params},
                                  headers={"origin": origin, "fly-client-ip": ip})

    async def test_a_name_is_checked_under_ten_endings_and_suggestions_skip_those(self):
        r = await self.ask("Tea Hub")
        body = r.json()
        self.assertEqual(r.status_code, 200)
        self.assertEqual([d["name"] for d in body["exact_domains"]], [f"teahub.{e}" for e in domains.ENDINGS])
        com, net = body["exact_domains"][:2]
        self.assertEqual((com["registrable"], com["registration_cost"], com["reason"]), (False, None, "domain_unavailable"))
        self.assertEqual((net["registrable"], net["registration_cost"], net["currency"]), (True, "10.46", "USD"))
        self.assertEqual([d["name"] for d in body["domains"]], ["teahubusa.com", "teahub.uk"])
        self.assertEqual((r.headers["access-control-allow-origin"], r.headers["vary"]), (ORIGIN, "Origin"))
        search = self.cf.calls[0]
        self.assertEqual((search.url.params["q"], search.url.params["limit"]), ("Tea Hub", "24"))

    async def test_a_domain_is_checked_as_given(self):
        body = (await self.ask("Tea-Hub.dev", limit=5)).json()
        self.assertEqual([d["name"] for d in body["exact_domains"]], ["tea-hub.dev"])
        self.assertEqual(self.cf.calls[0].url.params["limit"], "5")

    async def test_names_cloudflare_cannot_register_leave_exact_empty(self):
        body = (await self.ask("a.b")).json()
        self.assertEqual((body["exact_domains"], len(body["domains"])), ([], 3))

    async def test_only_teahub_may_read_the_answer(self):
        r = await self.ask("teahub", origin="https://elsewhere.test")
        self.assertEqual(r.status_code, 200)
        self.assertNotIn("access-control-allow-origin", r.headers)
        preflight = await self.web.options("/domains", headers={"origin": ORIGIN,
                                                                "access-control-request-method": "GET"})
        self.assertEqual((preflight.status_code, preflight.headers["access-control-allow-origin"],
                          preflight.headers["access-control-allow-methods"]), (204, ORIGIN, "GET"))

    async def test_bad_searches_never_reach_cloudflare(self):
        for q, params in (("", {}), ("-tea", {}), ("t" * 101, {}), ("tea/hub", {}), ("tea", {"limit": "0"}),
                          ("tea", {"limit": "51"}), ("tea", {"limit": "many"})):
            r = await self.ask(q, **params)
            self.assertEqual(r.status_code, 400, (q, params))
            self.assertEqual(r.headers["cache-control"], "no-store")
        self.assertEqual(self.cf.calls, [])

    async def test_answers_are_kept_and_a_failure_is_not(self):
        self.cf.fail = 500
        failed = await self.ask("teahub")
        self.assertEqual((failed.status_code, failed.json()["error"]), (502, "Cloudflare could not answer just now."))
        self.cf.fail = 0
        first, again = await self.ask("teahub"), await self.ask("TeaHub")
        self.assertEqual((first.status_code, again.json()), (200, first.json()))
        self.assertEqual(len(self.cf.calls), 4)  # the failed pair, then one search and one check

    async def test_each_client_and_everyone_are_held_to_a_few_a_minute(self):
        for n in range(domains.PER_CLIENT):
            self.assertEqual((await self.ask(f"tea{n}")).status_code, 200)
        limited = await self.ask("one-more")
        self.assertEqual((limited.status_code, limited.json()["error"]), (429, "Too many searches. Try again in a minute."))
        self.assertEqual((await self.ask("tea0")).status_code, 200)  # a kept answer costs nothing
        for n in range(domains.IN_ALL - domains.PER_CLIENT):
            self.assertEqual((await self.ask(f"other{n}", ip=f"198.51.100.{n}")).status_code, 200)
        self.assertEqual((await self.ask("fresh", ip="192.0.2.1")).status_code, 429)

    async def test_without_cloudflare_settings_it_says_so(self):
        with mock.patch.dict(os.environ, {"CLOUDFLARE_API_TOKEN": ""}):
            r = await self.ask("teahub")
        self.assertEqual((r.status_code, r.json()["error"]), (503, "Domain search is not set up here."))


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