devplacepy/tests/e2e/admin/gateway.py

# retoor <retoor@molodetz.nl>

import http.server
import json
import socket
import socketserver
import threading

import requests

from devplacepy.database import get_table, refresh_snapshot
from tests.conftest import BASE_URL, login_user


def _promote_to_admin(username: str) -> None:
    users = get_table("users")
    user = users.find_one(username=username)
    if user:
        users.update({"uid": user["uid"], "role": "Admin"}, ["uid"])


def _admin_api_key(username: str) -> str:
    refresh_snapshot()
    return get_table("users").find_one(username=username)["api_key"]


def _fake_models_upstream(model_ids):
    sock = socket.socket()
    sock.bind(("127.0.0.1", 0))
    port = sock.getsockname()[1]
    sock.close()
    body = json.dumps({"data": [{"id": m} for m in model_ids]}).encode()

    class Handler(http.server.BaseHTTPRequestHandler):
        def do_GET(self):
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(body)

        def log_message(self, *args):
            pass

    httpd = socketserver.TCPServer(("127.0.0.1", port), Handler)
    thread = threading.Thread(target=httpd.serve_forever, daemon=True)
    thread.start()
    return httpd, port


def test_model_form_swaps_target_field_by_provider(page, seeded_db):
    user = seeded_db["alice"]
    _promote_to_admin(user["username"])
    key = _admin_api_key(user["username"])
    auth = {"X-API-KEY": key}

    httpd, port = _fake_models_upstream(["vendor/known-a", "vendor/known-b"])
    listing_provider = f"e2elisting{port}"
    blind_provider = f"e2eblind{port}"
    try:
        requests.post(
            f"{BASE_URL}/admin/gateway/providers",
            json={
                "name": listing_provider,
                "base_url": f"http://127.0.0.1:{port}/v1/chat/completions",
            },
            headers=auth,
        )
        requests.post(
            f"{BASE_URL}/admin/gateway/providers",
            json={"name": blind_provider, "base_url": ""},
            headers=auth,
        )

        login_user(page, user)
        page.goto(f"{BASE_URL}/admin/gateway/models/new", wait_until="domcontentloaded")

        target_input = page.locator("#gw-model-target")
        target_select = page.locator("#gw-model-target-select")
        target_input.wait_for(state="visible")
        assert not target_select.is_visible()

        page.select_option("#gw-model-provider", listing_provider)
        target_select.wait_for(state="visible", timeout=10000)
        assert not target_input.is_visible()
        assert target_select.get_attribute("required") is not None
        options = target_select.locator("option").all_inner_texts()
        assert "vendor/known-a" in options
        assert "vendor/known-b" in options

        page.select_option("#gw-model-provider", blind_provider)
        target_input.wait_for(state="visible", timeout=10000)
        assert not target_select.is_visible()
    finally:
        httpd.shutdown()
        requests.delete(
            f"{BASE_URL}/admin/gateway/providers/{listing_provider}", headers=auth
        )
        requests.delete(
            f"{BASE_URL}/admin/gateway/providers/{blind_provider}", headers=auth
        )


def test_model_form_submits_selected_model_from_dropdown(page, seeded_db):
    user = seeded_db["alice"]
    _promote_to_admin(user["username"])
    key = _admin_api_key(user["username"])
    auth = {"X-API-KEY": key}

    httpd, port = _fake_models_upstream(["vendor/pick-me"])
    provider = f"e2esubmit{port}"
    source = f"e2esource{port}"
    try:
        requests.post(
            f"{BASE_URL}/admin/gateway/providers",
            json={
                "name": provider,
                "base_url": f"http://127.0.0.1:{port}/v1/chat/completions",
            },
            headers=auth,
        )

        login_user(page, user)
        page.goto(f"{BASE_URL}/admin/gateway/models/new", wait_until="domcontentloaded")
        page.fill("#gw-model-source", source)
        page.select_option("#gw-model-provider", provider)
        page.locator("#gw-model-target-select").wait_for(state="visible", timeout=10000)
        page.select_option("#gw-model-target-select", "vendor/pick-me")
        page.click("button[type='submit']")
        page.wait_for_url(f"{BASE_URL}/admin/gateway?tab=models", wait_until="domcontentloaded")

        listed = requests.get(f"{BASE_URL}/admin/gateway/models", headers=auth).json()
        row = next(m for m in listed["models"] if m["source_model"] == source)
        assert row["target_model"] == "vendor/pick-me"
        assert row["provider"] == provider
    finally:
        httpd.shutdown()
        requests.delete(f"{BASE_URL}/admin/gateway/models/{source}", headers=auth)
        requests.delete(f"{BASE_URL}/admin/gateway/providers/{provider}", headers=auth)