devplacepy/tests/e2e/openai/v1/chat/completions.py

# retoor <retoor@molodetz.nl>

import json
import requests
from starlette.requests import Request
from tests.conftest import BASE_URL, run_async
from devplacepy.database import get_table, set_setting
from devplacepy.utils import generate_uid
import devplacepy.services.openai_gateway.gateway as gwmod
from devplacepy.services.openai_gateway import GatewayService
class FakeResp_openai_gateway:
    def __init__(self, status=200, payload=None, ctype="application/json", content=b""):
        self.status_code = status
        self._payload = payload
        self.text = json.dumps(payload) if payload is not None else ""
        self.headers = {"content-type": ctype}
        self.content = content

    def json(self):
        if self._payload is None:
            raise ValueError("no json")
        return self._payload
class FakeRequest:
    def __init__(self, method, url, json_body):
        self.method = method
        self.url = url
        self.json_body = json_body
        self.extensions = {}
class FakeClient_openai_gateway:
    def __init__(self, *a, **k):
        self.calls = []

    def build_request(self, method, url, headers=None, json=None, content=None):
        return FakeRequest(method, url, json)

    async def send(self, request):
        self.calls.append((request.url, request.json_body))
        body = request.json_body or {}
        return FakeResp_openai_gateway(
            payload={
                "id": "x",
                "model": body.get("model"),
                "choices": [{"message": {"content": "hi there"}}],
            }
        )

    async def post(self, url, headers=None, json=None, timeout=None):
        self.calls.append((url, json))
        return FakeResp_openai_gateway(
            payload={
                "id": "x",
                "model": json.get("model"),
                "choices": [{"message": {"content": "hi there"}}],
            }
        )

    async def request(self, method, url, headers=None, content=None):
        return FakeResp_openai_gateway(payload={"ok": True})

    async def aclose(self):
        pass
def _make_request_openai_gateway(headers=None, cookies=None):
    headers = headers or {}
    raw = [(k.lower().encode(), v.encode()) for k, v in headers.items()]
    if cookies:
        raw.append(
            (b"cookie", "; ".join(f"{k}={v}" for k, v in cookies.items()).encode())
        )
    return Request(
        {
            "type": "http",
            "method": "POST",
            "path": "/openai/v1/chat/completions",
            "query_string": b"",
            "headers": raw,
            "state": {},
        }
    )
def _make_admin_openai_gateway(role="Admin"):
    username = f"gw_{generate_uid()[:8]}"
    api_key = generate_uid()
    get_table("users").insert(
        {
            "uid": generate_uid(),
            "username": username,
            "terms_version": "1",
            "email": f"{username}@t.dev",
            "api_key": api_key,
            "role": role,
            "is_active": True,
        }
    )
    return username, api_key
def _config(page, **fields):
    page.request.post(f"{BASE_URL}/admin/services/openai/config", form=fields)


def _grant_ai_consent(page, username):
    page.request.post(
        f"{BASE_URL}/profile/{username}/consent",
        form={"kind": "ai_third_party", "granted": "1"},
    )


def test_gateway_disabled_returns_503(alice):
    page, _ = alice
    page.request.post(f"{BASE_URL}/admin/services/openai/stop")
    try:
        r = requests.post(
            f"{BASE_URL}/openai/v1/chat/completions", json={"messages": []}
        )
        assert r.status_code == 503
    finally:
        page.request.post(f"{BASE_URL}/admin/services/openai/start")


def test_gateway_auth_and_routing(alice):
    page, user = alice
    page.request.post(f"{BASE_URL}/admin/services/openai/start")
    _config(
        page,
        gateway_upstream_url="http://127.0.0.1:9/chat/completions",
        gateway_vision_enabled="0",
        gateway_allow_admins="1",
        gateway_access_key="gwkey",
    )
    try:
        no_creds = requests.post(
            f"{BASE_URL}/openai/v1/chat/completions", json={"messages": []}
        )
        assert no_creds.status_code == 401

        with_key = requests.post(
            f"{BASE_URL}/openai/v1/chat/completions",
            headers={"X-API-KEY": "gwkey"},
            json={"messages": [{"role": "user", "content": "hi"}]},
        )
        assert with_key.status_code == 502  # auth passed, upstream unreachable

        admin_key = get_table("users").find_one(username=user["username"])["api_key"]
        no_consent = requests.post(
            f"{BASE_URL}/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {admin_key}"},
            json={"messages": [{"role": "user", "content": "hi"}]},
        )
        assert no_consent.status_code == 403, (
            "a user's own key must not reach the provider without ai_third_party consent"
        )

        _grant_ai_consent(page, user["username"])
        with_admin = requests.post(
            f"{BASE_URL}/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {admin_key}"},
            json={"messages": [{"role": "user", "content": "hi"}]},
        )
        assert with_admin.status_code == 502
    finally:
        _config(page, gateway_access_key="")
        page.request.post(f"{BASE_URL}/admin/services/openai/start")