mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 22:14:46 +00:00
feat: BYOM (#2433)
This commit is contained in:
68 files changed
+7645
-282
No files matched your search
@@ -0,0 +1,688 @@
|
||||
"""Tests for the BYOM REST API at /api/user/models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
"""Patch the routes' db helpers to yield the given pg connection."""
|
||||
|
||||
@contextmanager
|
||||
def _yield_conn():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.user.models.routes.db_session", _yield_conn
|
||||
), patch(
|
||||
"application.api.user.models.routes.db_readonly", _yield_conn
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_registry():
|
||||
from application.core.model_registry import ModelRegistry
|
||||
|
||||
ModelRegistry.reset()
|
||||
yield
|
||||
ModelRegistry.reset()
|
||||
|
||||
|
||||
# Auth
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAuth:
|
||||
def test_list_unauthenticated_returns_401(self, app):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with app.test_request_context("/api/user/models"):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = None
|
||||
resp = UserModelsCollectionResource().get()
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_create_unauthenticated_returns_401(self, app):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "x",
|
||||
"display_name": "x",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"api_key": "k",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = None
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
# Create
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCreate:
|
||||
def test_creates_and_returns_201_without_api_key(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
# Mock DNS so the SSRF check passes for api.mistral.ai without
|
||||
# hitting the network.
|
||||
with patch("application.security.safe_url.socket.getaddrinfo") as gai:
|
||||
gai.return_value = [
|
||||
(None, None, None, None, ("104.18.0.1", 0))
|
||||
]
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "mistral-large-latest",
|
||||
"display_name": "My Mistral",
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"api_key": "sk-mistral-test",
|
||||
"capabilities": {
|
||||
"supports_tools": True,
|
||||
"context_window": 128000,
|
||||
},
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
|
||||
assert resp.status_code == 201
|
||||
body = resp.get_json()
|
||||
assert body["upstream_model_id"] == "mistral-large-latest"
|
||||
assert body["source"] == "user"
|
||||
# Critical: api_key must NEVER appear in the response
|
||||
assert "api_key" not in body
|
||||
for v in body.values():
|
||||
assert v != "sk-mistral-test"
|
||||
|
||||
def test_create_rejects_missing_required_fields(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={"upstream_model_id": "x"}, # missing the others
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_create_rejects_loopback_url(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "x",
|
||||
"display_name": "x",
|
||||
"base_url": "https://127.0.0.1/v1",
|
||||
"api_key": "k",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 400
|
||||
body = resp.get_json()
|
||||
assert "error" in body
|
||||
|
||||
def test_create_rejects_unknown_attachment_alias(self, app, pg_conn):
|
||||
"""The UI sends ``["image"]`` as an alias; bad strings ("video",
|
||||
typos) must reject at the boundary so the DB never holds
|
||||
garbage that the registry would later silently drop.
|
||||
"""
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with patch("application.security.safe_url.socket.getaddrinfo") as gai:
|
||||
gai.return_value = [
|
||||
(None, None, None, None, ("104.18.0.1", 0))
|
||||
]
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "m",
|
||||
"display_name": "M",
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"api_key": "k",
|
||||
"capabilities": {"attachments": ["video"]},
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 400
|
||||
body = resp.get_json()
|
||||
assert "video" in body["error"]
|
||||
|
||||
def test_create_accepts_image_alias_and_raw_mime(self, app, pg_conn):
|
||||
"""The known ``image`` alias and raw MIME types both pass."""
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with patch("application.security.safe_url.socket.getaddrinfo") as gai:
|
||||
gai.return_value = [
|
||||
(None, None, None, None, ("104.18.0.1", 0))
|
||||
]
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "m",
|
||||
"display_name": "M",
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"api_key": "k",
|
||||
"capabilities": {
|
||||
"attachments": ["image", "application/pdf"],
|
||||
},
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 201
|
||||
|
||||
def test_create_rejects_private_ip_dns(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with patch("application.security.safe_url.socket.getaddrinfo") as gai:
|
||||
# Hostname resolves to a private IP only — must reject
|
||||
gai.return_value = [
|
||||
(None, None, None, None, ("10.0.0.5", 0))
|
||||
]
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "x",
|
||||
"display_name": "x",
|
||||
"base_url": "https://evil.example.com/v1",
|
||||
"api_key": "k",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
# List / get / patch / delete
|
||||
|
||||
|
||||
def _create_via_repo(pg_conn, user_id="user-1", **kwargs):
|
||||
from application.storage.db.repositories.user_custom_models import (
|
||||
UserCustomModelsRepository,
|
||||
)
|
||||
|
||||
return UserCustomModelsRepository(pg_conn).create(
|
||||
user_id=user_id,
|
||||
upstream_model_id=kwargs.pop("upstream_model_id", "mistral-large-latest"),
|
||||
display_name=kwargs.pop("display_name", "My Mistral"),
|
||||
base_url=kwargs.pop("base_url", "https://api.mistral.ai/v1"),
|
||||
api_key_plaintext=kwargs.pop("api_key_plaintext", "sk-mistral-test"),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestList:
|
||||
def test_lists_only_users_own(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
_create_via_repo(pg_conn, user_id="alice", upstream_model_id="alice-1")
|
||||
_create_via_repo(pg_conn, user_id="bob", upstream_model_id="bob-1")
|
||||
|
||||
with app.test_request_context("/api/user/models"):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().get()
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
upstream_ids = {m["upstream_model_id"] for m in body["models"]}
|
||||
assert upstream_ids == {"alice-1"}
|
||||
# Never expose the api_key
|
||||
for m in body["models"]:
|
||||
assert "api_key" not in m
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGet:
|
||||
def test_returns_404_for_other_users_model(self, app, pg_conn):
|
||||
from application.api.user.models.routes import UserModelResource
|
||||
|
||||
created = _create_via_repo(pg_conn, user_id="alice")
|
||||
with app.test_request_context(
|
||||
f"/api/user/models/{created['id']}"
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "bob"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelResource().get(model_id=created["id"])
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPatch:
|
||||
def test_patch_updates_display_name(self, app, pg_conn):
|
||||
from application.api.user.models.routes import UserModelResource
|
||||
|
||||
created = _create_via_repo(pg_conn, user_id="user-1")
|
||||
with app.test_request_context(
|
||||
f"/api/user/models/{created['id']}",
|
||||
method="PATCH",
|
||||
json={"display_name": "Renamed"},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelResource().patch(model_id=created["id"])
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
assert body["display_name"] == "Renamed"
|
||||
|
||||
def test_patch_blank_api_key_keeps_existing(self, app, pg_conn):
|
||||
"""Critical PATCH semantic: empty/missing api_key in body must
|
||||
preserve the stored ciphertext (the UI sends a blank password
|
||||
field when the user wants to keep the existing key)."""
|
||||
from application.api.user.models.routes import UserModelResource
|
||||
from application.storage.db.repositories.user_custom_models import (
|
||||
UserCustomModelsRepository,
|
||||
)
|
||||
|
||||
created = _create_via_repo(pg_conn, user_id="user-1")
|
||||
original_key_plaintext = "sk-mistral-test"
|
||||
|
||||
with app.test_request_context(
|
||||
f"/api/user/models/{created['id']}",
|
||||
method="PATCH",
|
||||
json={"display_name": "Just rename me", "api_key": ""},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelResource().patch(model_id=created["id"])
|
||||
assert resp.status_code == 200
|
||||
|
||||
# Decrypted key is unchanged
|
||||
repo = UserCustomModelsRepository(pg_conn)
|
||||
assert (
|
||||
repo.get_decrypted_api_key(created["id"], "user-1")
|
||||
== original_key_plaintext
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDelete:
|
||||
def test_delete_removes_row_and_invalidates_cache(self, app, pg_conn):
|
||||
from application.api.user.models.routes import UserModelResource
|
||||
from application.core.model_registry import ModelRegistry
|
||||
|
||||
created = _create_via_repo(pg_conn, user_id="user-1")
|
||||
# Warm the registry's per-user cache via a lookup
|
||||
with patch(
|
||||
"application.storage.db.session.db_readonly"
|
||||
) as ro:
|
||||
@contextmanager
|
||||
def _y():
|
||||
yield pg_conn
|
||||
|
||||
ro.side_effect = _y
|
||||
ModelRegistry.get_instance().get_model(created["id"], user_id="user-1")
|
||||
# Now delete via the route — invalidation must happen so a
|
||||
# subsequent lookup misses
|
||||
with app.test_request_context(
|
||||
f"/api/user/models/{created['id']}", method="DELETE"
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelResource().delete(model_id=created["id"])
|
||||
assert resp.status_code == 200
|
||||
# Cache invalidated → next lookup re-queries DB and finds nothing
|
||||
assert "user-1" not in ModelRegistry.get_instance()._user_models
|
||||
|
||||
|
||||
# /api/models combined view
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSecurityCreateRejectsBlankFields:
|
||||
"""P1 #1 partial: blank api_key on create must be rejected so we
|
||||
can never end up with an unroutable BYOM record that would cause
|
||||
LLMCreator to leak settings.API_KEY to the user-supplied URL."""
|
||||
|
||||
def test_create_rejects_blank_api_key(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelsCollectionResource,
|
||||
)
|
||||
|
||||
with patch("application.security.safe_url.socket.getaddrinfo") as gai:
|
||||
gai.return_value = [(None, None, None, None, ("104.18.0.1", 0))]
|
||||
with app.test_request_context(
|
||||
"/api/user/models",
|
||||
method="POST",
|
||||
json={
|
||||
"upstream_model_id": "x",
|
||||
"display_name": "x",
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"api_key": " ", # whitespace only
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "u"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelsCollectionResource().post()
|
||||
assert resp.status_code == 400
|
||||
body = resp.get_json()
|
||||
assert "api_key" in (body.get("error") or "").lower()
|
||||
|
||||
def test_patch_rejects_blank_required_field(self, app, pg_conn):
|
||||
from application.api.user.models.routes import UserModelResource
|
||||
from application.storage.db.repositories.user_custom_models import (
|
||||
UserCustomModelsRepository,
|
||||
)
|
||||
|
||||
created = UserCustomModelsRepository(pg_conn).create(
|
||||
user_id="u",
|
||||
upstream_model_id="x",
|
||||
display_name="x",
|
||||
base_url="https://api.mistral.ai/v1",
|
||||
api_key_plaintext="sk-x",
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
f"/api/user/models/{created['id']}",
|
||||
method="PATCH",
|
||||
json={"base_url": " "},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "u"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelResource().patch(model_id=created["id"])
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPayloadConnectionTest:
|
||||
"""Verifies the payload-based test endpoint. Lets the UI's 'Test
|
||||
connection' button work *before* the model is saved — operators
|
||||
expect to validate their endpoint + key before committing."""
|
||||
|
||||
def test_payload_test_rejects_unsafe_url(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelTestPayloadResource,
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/user/models/test",
|
||||
method="POST",
|
||||
json={
|
||||
"base_url": "https://127.0.0.1/v1",
|
||||
"api_key": "sk-anything",
|
||||
"upstream_model_id": "x",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelTestPayloadResource().post()
|
||||
assert resp.status_code == 400
|
||||
body = resp.get_json()
|
||||
assert body["ok"] is False
|
||||
|
||||
def test_payload_test_returns_ok_when_upstream_responds_2xx(
|
||||
self, app, pg_conn
|
||||
):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelTestPayloadResource,
|
||||
)
|
||||
|
||||
# pinned_post is the IP-pinned dispatch helper. Patching it
|
||||
# bypasses both the SSRF guard and the network — the success
|
||||
# path we're verifying here is the route's response handling.
|
||||
with patch("application.api.user.models.routes.pinned_post") as rp:
|
||||
rp.return_value = MagicMock(
|
||||
status_code=200,
|
||||
headers={"Content-Type": "application/json"},
|
||||
text='{"ok": true}',
|
||||
)
|
||||
with app.test_request_context(
|
||||
"/api/user/models/test",
|
||||
method="POST",
|
||||
json={
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"api_key": "sk-mistral-test",
|
||||
"upstream_model_id": "mistral-large-latest",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelTestPayloadResource().post()
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
assert body["ok"] is True
|
||||
# Verify the upstream call carried the user's submitted key (not
|
||||
# whatever's in the DB) and the right model name.
|
||||
call_args = rp.call_args
|
||||
assert call_args.kwargs["headers"]["Authorization"] == "Bearer sk-mistral-test"
|
||||
assert call_args.kwargs["json"]["model"] == "mistral-large-latest"
|
||||
|
||||
def test_payload_test_unauthenticated_returns_401(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelTestPayloadResource,
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/user/models/test",
|
||||
method="POST",
|
||||
json={
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"api_key": "k",
|
||||
"upstream_model_id": "x",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = None
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelTestPayloadResource().post()
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_payload_test_missing_fields_returns_400(self, app, pg_conn):
|
||||
from application.api.user.models.routes import (
|
||||
UserModelTestPayloadResource,
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/user/models/test",
|
||||
method="POST",
|
||||
json={"base_url": "https://api.mistral.ai/v1"},
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
with _patch_db(pg_conn):
|
||||
resp = UserModelTestPayloadResource().post()
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestByIdConnectionTestAcceptsOverrides:
|
||||
"""P3: in edit mode the modal sends current form state as overrides
|
||||
so the test reflects in-flight edits (not the saved record)."""
|
||||
|
||||
def _make_row(self, pg_conn):
|
||||
from application.storage.db.repositories.user_custom_models import (
|
||||
UserCustomModelsRepository,
|
||||
)
|
||||
|
||||
return UserCustomModelsRepository(pg_conn).create(
|
||||
user_id="u",
|
||||
upstream_model_id="stored-model",
|
||||
display_name="Stored",
|
||||
base_url="https://stored.example.com/v1",
|
||||
api_key_plaintext="sk-stored",
|
||||
)
|
||||
|
||||
def _post_test(self, app, pg_conn, model_id, body):
|
||||
from application.api.user.models.routes import UserModelTestResource
|
||||
|
||||
with patch("application.api.user.models.routes.pinned_post") as rp:
|
||||
rp.return_value = MagicMock(
|
||||
status_code=200,
|
||||
headers={"Content-Type": "application/json"},
|
||||
text='{"ok": true}',
|
||||
)
|
||||
with app.test_request_context(
|
||||
f"/api/user/models/{model_id}/test", method="POST", json=body
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "u"}
|
||||
with _patch_db(pg_conn):
|
||||
UserModelTestResource().post(model_id=model_id)
|
||||
return rp.call_args
|
||||
|
||||
def test_overrides_win_when_supplied(self, app, pg_conn):
|
||||
row = self._make_row(pg_conn)
|
||||
ca = self._post_test(
|
||||
app,
|
||||
pg_conn,
|
||||
row["id"],
|
||||
{
|
||||
"base_url": "https://new.example.com/v1",
|
||||
"api_key": "sk-new",
|
||||
"upstream_model_id": "new-model",
|
||||
},
|
||||
)
|
||||
assert ca.args[0] == "https://new.example.com/v1/chat/completions"
|
||||
assert ca.kwargs["headers"]["Authorization"] == "Bearer sk-new"
|
||||
assert ca.kwargs["json"]["model"] == "new-model"
|
||||
|
||||
def test_blank_overrides_fall_back_to_stored(self, app, pg_conn):
|
||||
"""The classic edit-mode flow: user changed base_url, left
|
||||
api_key blank — server uses the new URL but the stored key."""
|
||||
row = self._make_row(pg_conn)
|
||||
ca = self._post_test(
|
||||
app,
|
||||
pg_conn,
|
||||
row["id"],
|
||||
{
|
||||
"base_url": "https://new.example.com/v1",
|
||||
"api_key": "",
|
||||
"upstream_model_id": "",
|
||||
},
|
||||
)
|
||||
assert ca.args[0] == "https://new.example.com/v1/chat/completions"
|
||||
# Stored key (decrypted) was used.
|
||||
assert ca.kwargs["headers"]["Authorization"] == "Bearer sk-stored"
|
||||
# Stored upstream_model_id was used.
|
||||
assert ca.kwargs["json"]["model"] == "stored-model"
|
||||
|
||||
def test_empty_body_uses_all_stored_values(self, app, pg_conn):
|
||||
row = self._make_row(pg_conn)
|
||||
ca = self._post_test(app, pg_conn, row["id"], {})
|
||||
assert ca.args[0] == "https://stored.example.com/v1/chat/completions"
|
||||
assert ca.kwargs["headers"]["Authorization"] == "Bearer sk-stored"
|
||||
assert ca.kwargs["json"]["model"] == "stored-model"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestApiModelsListWithUser:
|
||||
def test_includes_user_models_when_authenticated(self, app, pg_conn):
|
||||
"""GET /api/models with auth should surface the user's BYOM
|
||||
records alongside built-ins, each tagged with `source`."""
|
||||
from application.api.user.models.routes import ModelsListResource
|
||||
from application.core.model_registry import ModelRegistry
|
||||
|
||||
created = _create_via_repo(
|
||||
pg_conn, user_id="user-1", display_name="My Mistral"
|
||||
)
|
||||
|
||||
# Patch the *registry's* db_readonly so the per-user layer load
|
||||
# uses the test connection.
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with patch(
|
||||
"application.storage.db.session.db_readonly", _yield
|
||||
):
|
||||
ModelRegistry.reset()
|
||||
with app.test_request_context("/api/models"):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "user-1"}
|
||||
resp = ModelsListResource().get()
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
ids = [m["id"] for m in body["models"]]
|
||||
assert created["id"] in ids
|
||||
# Source label tags it for the UI
|
||||
user_entries = [m for m in body["models"] if m["id"] == created["id"]]
|
||||
assert user_entries[0]["source"] == "user"
|
||||
# Built-ins still present
|
||||
assert any(m.get("source") == "builtin" for m in body["models"])
|
||||
Reference in new issue
Block a user