mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 02:13:47 +00:00
Pg 4 (#2390)
* feat: postgres tests * feat: mongo cutoff * feat: mongo cutoff * feat: adjust docs and compose files * fix: mini code mongo removals * fix: tests and k8s mongo stuff * feat: test fixes * fix: ruff * fix: vale * Potential fix for pull request finding 'CodeQL / Clear-text logging of sensitive information' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> * fix: mini suggestions * vale lint fix 2 * fix: codeql columns thing * fix: test mongo * fix: tests coverage * feat: better tests 4 * feat: more tests * feat: decent coverage * fix: ruff fixes * fix: remove mongo mock * feat: enhance workflow engine and API routes; add document retrieval and source handling * feat: e2e tests * fix: mcp, mongo and more * fix: mini codeql warning * fix: agent chunk view * fix: mini issues * fix: more pg fixes * feat: postgres prep on start * feat: qa tests * fix: mini improvements * fix: tests --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Co-authored-by: Siddhant Rai <siddhant.rai.5686@gmail.com>
This commit is contained in:
312 files changed
+763725
-25159
No files matched your search
+285
-576
@@ -1,611 +1,320 @@
|
||||
"""Unit tests for application/api/internal/routes.py.
|
||||
"""Tests for application/api/internal/routes.py.
|
||||
|
||||
Covers:
|
||||
- verify_internal_key: key validation
|
||||
- /api/download: file download
|
||||
- /api/upload_index: index file upload (existing & new entries)
|
||||
Uses the ephemeral ``pg_conn`` fixture so the sources repository writes
|
||||
happen against a real Postgres schema.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from bson.objectid import ObjectId
|
||||
from flask import Flask
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def internal_app(monkeypatch, mock_mongo_db):
|
||||
"""Create a Flask app with the internal blueprint registered."""
|
||||
from flask import Flask
|
||||
_TEST_KEY = "test-internal-key"
|
||||
_AUTH = {"X-Internal-Key": _TEST_KEY}
|
||||
|
||||
# Patch module-level MongoDB references before importing routes
|
||||
from application.core.settings import settings
|
||||
|
||||
db = mock_mongo_db[settings.MONGO_DB_NAME]
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.conversations_collection",
|
||||
db["conversations"],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.sources_collection",
|
||||
db["sources"],
|
||||
)
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.internal.routes.db_session", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _make_app():
|
||||
from application.api.internal.routes import internal
|
||||
|
||||
app = Flask(__name__)
|
||||
app.register_blueprint(internal)
|
||||
app.config["TESTING"] = True
|
||||
return app, db
|
||||
return app
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# verify_internal_key
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestVerifyInternalKey:
|
||||
def test_rejects_when_internal_key_not_configured(self):
|
||||
app = _make_app()
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", ""
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.get("/api/download")
|
||||
assert r.status_code == 401
|
||||
|
||||
def test_no_internal_key_configured_rejects_access(
|
||||
self, internal_app, monkeypatch
|
||||
):
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings",
|
||||
MagicMock(
|
||||
INTERNAL_KEY=None,
|
||||
UPLOAD_FOLDER="uploads",
|
||||
VECTOR_STORE="faiss",
|
||||
EMBEDDINGS_NAME="test",
|
||||
MONGO_DB_NAME="docsgpt",
|
||||
),
|
||||
)
|
||||
with app.test_client() as client:
|
||||
resp = client.get("/api/download?user=u&name=n&file=f")
|
||||
assert resp.status_code == 401
|
||||
def test_rejects_when_key_missing(self):
|
||||
app = _make_app()
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.get("/api/download")
|
||||
assert r.status_code == 401
|
||||
|
||||
def test_missing_key_returns_401(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings",
|
||||
MagicMock(
|
||||
INTERNAL_KEY="secret123",
|
||||
UPLOAD_FOLDER="uploads",
|
||||
VECTOR_STORE="faiss",
|
||||
EMBEDDINGS_NAME="test",
|
||||
MONGO_DB_NAME="docsgpt",
|
||||
),
|
||||
)
|
||||
with app.test_client() as client:
|
||||
resp = client.get("/api/download?user=u&name=n&file=f")
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_wrong_key_returns_401(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings",
|
||||
MagicMock(
|
||||
INTERNAL_KEY="secret123",
|
||||
UPLOAD_FOLDER="uploads",
|
||||
VECTOR_STORE="faiss",
|
||||
EMBEDDINGS_NAME="test",
|
||||
MONGO_DB_NAME="docsgpt",
|
||||
),
|
||||
)
|
||||
with app.test_client() as client:
|
||||
resp = client.get(
|
||||
"/api/download?user=u&name=n&file=f",
|
||||
headers={"X-Internal-Key": "wrong"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_correct_key_allows_access(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings",
|
||||
MagicMock(
|
||||
INTERNAL_KEY="secret123",
|
||||
UPLOAD_FOLDER="uploads",
|
||||
VECTOR_STORE="faiss",
|
||||
EMBEDDINGS_NAME="test",
|
||||
MONGO_DB_NAME="docsgpt",
|
||||
),
|
||||
)
|
||||
with app.test_client() as client:
|
||||
# Will 404 for missing file, but should pass auth check
|
||||
resp = client.get(
|
||||
"/api/download?user=u&name=n&file=f",
|
||||
headers={"X-Internal-Key": "secret123"},
|
||||
)
|
||||
assert resp.status_code != 401
|
||||
def test_rejects_when_key_mismatch(self):
|
||||
app = _make_app()
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.get("/api/download", headers={"X-Internal-Key": "wrong"})
|
||||
assert r.status_code == 401
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /api/upload_index
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestDownloadFile:
|
||||
def test_returns_404_for_missing_file(self, tmp_path):
|
||||
app = _make_app()
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.UPLOAD_FOLDER",
|
||||
str(tmp_path),
|
||||
), patch(
|
||||
"application.api.internal.routes.current_dir", ""
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.get(
|
||||
"/api/download?user=alice&name=job1&file=missing.txt",
|
||||
headers=_AUTH,
|
||||
)
|
||||
assert r.status_code == 404
|
||||
|
||||
def test_returns_file_when_present(self, tmp_path):
|
||||
app = _make_app()
|
||||
user_dir = tmp_path / "bob" / "job1"
|
||||
user_dir.mkdir(parents=True)
|
||||
(user_dir / "hello.txt").write_text("hi there")
|
||||
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.UPLOAD_FOLDER",
|
||||
str(tmp_path),
|
||||
), patch(
|
||||
"application.api.internal.routes.current_dir", ""
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.get(
|
||||
"/api/download?user=bob&name=job1&file=hello.txt",
|
||||
headers=_AUTH,
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert r.data == b"hi there"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestUploadIndex:
|
||||
def _base_form(self, *, source_id="source-1"):
|
||||
return {
|
||||
"user": "alice",
|
||||
"name": "Job A",
|
||||
"tokens": "100",
|
||||
"retriever": "classic",
|
||||
"id": source_id,
|
||||
"type": "file",
|
||||
}
|
||||
|
||||
_TEST_INTERNAL_KEY = "test-internal-key"
|
||||
_AUTH_HEADERS = {"X-Internal-Key": "test-internal-key"}
|
||||
def test_rejects_without_auth(self):
|
||||
app = _make_app()
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index")
|
||||
assert r.status_code == 401
|
||||
|
||||
def _make_settings(self, vector_store="faiss"):
|
||||
return MagicMock(
|
||||
INTERNAL_KEY=self._TEST_INTERNAL_KEY,
|
||||
UPLOAD_FOLDER="uploads",
|
||||
VECTOR_STORE=vector_store,
|
||||
EMBEDDINGS_NAME="test_embeddings",
|
||||
MONGO_DB_NAME="docsgpt",
|
||||
def test_rejects_missing_user(self):
|
||||
app = _make_app()
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index", headers=_AUTH, data={})
|
||||
assert r.status_code == 200
|
||||
assert r.json == {"status": "no user"}
|
||||
|
||||
def test_rejects_missing_name(self):
|
||||
app = _make_app()
|
||||
form = {"user": "alice"}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index", headers=_AUTH, data=form)
|
||||
assert r.status_code == 200
|
||||
assert r.json == {"status": "no name"}
|
||||
|
||||
def test_creates_new_source_for_non_faiss_store(self, pg_conn):
|
||||
"""For non-faiss VECTOR_STORE the route skips file uploads entirely."""
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
app = _make_app()
|
||||
form = {**self._base_form(source_id="legacy-source-1")}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "milvus"
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.EMBEDDINGS_NAME", "emb"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=MagicMock(),
|
||||
), _patch_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index", headers=_AUTH, data=form)
|
||||
|
||||
assert r.status_code == 200
|
||||
assert r.json == {"status": "ok"}
|
||||
repo = SourcesRepository(pg_conn)
|
||||
found = repo.get_by_legacy_id("legacy-source-1", "alice")
|
||||
assert found is not None
|
||||
|
||||
def test_updates_existing_source(self, pg_conn):
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
repo = SourcesRepository(pg_conn)
|
||||
created = repo.create(
|
||||
"initial",
|
||||
user_id="alice",
|
||||
legacy_mongo_id="legacy-src-2",
|
||||
tokens="0",
|
||||
)
|
||||
app = _make_app()
|
||||
form = {**self._base_form(source_id="legacy-src-2"), "tokens": "999"}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "milvus"
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.EMBEDDINGS_NAME", "emb"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=MagicMock(),
|
||||
), _patch_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index", headers=_AUTH, data=form)
|
||||
assert r.status_code == 200
|
||||
|
||||
def test_missing_user_returns_no_user(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", self._make_settings()
|
||||
)
|
||||
with app.test_client() as client:
|
||||
resp = client.post("/api/upload_index", data={}, headers=self._AUTH_HEADERS)
|
||||
assert resp.json["status"] == "no user"
|
||||
updated = repo.get(str(created["id"]), "alice")
|
||||
assert updated["tokens"] == "999"
|
||||
|
||||
def test_missing_name_returns_no_name(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", self._make_settings()
|
||||
)
|
||||
with app.test_client() as client:
|
||||
resp = client.post("/api/upload_index", data={"user": "testuser"}, headers=self._AUTH_HEADERS)
|
||||
assert resp.json["status"] == "no name"
|
||||
def test_handles_directory_structure_and_file_name_map(self, pg_conn):
|
||||
import json
|
||||
|
||||
def test_creates_new_source_entry(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
app = _make_app()
|
||||
form = {
|
||||
**self._base_form(source_id="legacy-src-3"),
|
||||
"directory_structure": json.dumps({"root": {"a.txt": None}}),
|
||||
"file_name_map": json.dumps({"a.txt": "Original A.txt"}),
|
||||
}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "milvus"
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.EMBEDDINGS_NAME", "emb"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=MagicMock(),
|
||||
), _patch_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index", headers=_AUTH, data=form)
|
||||
assert r.status_code == 200
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "testuser",
|
||||
"name": "testjob",
|
||||
"tokens": "100",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
def test_invalid_json_falls_back_to_empty(self, pg_conn):
|
||||
app = _make_app()
|
||||
form = {
|
||||
**self._base_form(source_id="legacy-src-4"),
|
||||
"directory_structure": "not-json",
|
||||
"file_name_map": "also-not-json",
|
||||
}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "milvus"
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.EMBEDDINGS_NAME", "emb"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=MagicMock(),
|
||||
), _patch_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
r = c.post("/api/upload_index", headers=_AUTH, data=form)
|
||||
assert r.status_code == 200
|
||||
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert entry is not None
|
||||
assert entry["user"] == "testuser"
|
||||
assert entry["name"] == "testjob"
|
||||
def test_faiss_missing_file_faiss(self, pg_conn):
|
||||
app = _make_app()
|
||||
form = self._base_form(source_id="legacy-src-5")
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "faiss"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.post(
|
||||
"/api/upload_index",
|
||||
headers=_AUTH,
|
||||
data=form,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert r.json == {"status": "no file"}
|
||||
|
||||
def test_updates_existing_source_entry(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
doc_id = ObjectId()
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
def test_faiss_missing_file_pkl(self, pg_conn):
|
||||
app = _make_app()
|
||||
data = {
|
||||
**self._base_form(source_id="legacy-src-6"),
|
||||
"file_faiss": (io.BytesIO(b"faiss-data"), "index.faiss"),
|
||||
}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "faiss"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
with app.test_client() as c:
|
||||
r = c.post(
|
||||
"/api/upload_index",
|
||||
headers=_AUTH,
|
||||
data=data,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert r.json == {"status": "no file"}
|
||||
|
||||
# Insert existing entry
|
||||
db["sources"].insert_one(
|
||||
{"_id": doc_id, "user": "old_user", "name": "old_name"}
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "new_user",
|
||||
"name": "new_name",
|
||||
"tokens": "200",
|
||||
"retriever": "hybrid",
|
||||
"id": str(doc_id),
|
||||
"type": "remote",
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": doc_id})
|
||||
assert entry["user"] == "new_user"
|
||||
assert entry["name"] == "new_name"
|
||||
|
||||
def test_parses_directory_structure_json(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
dir_struct = {"root": {"files": ["a.txt", "b.txt"]}}
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"directory_structure": json.dumps(dir_struct),
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert entry["directory_structure"] == dir_struct
|
||||
|
||||
def test_invalid_directory_structure_defaults_empty(
|
||||
self, internal_app, monkeypatch
|
||||
):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"directory_structure": "not valid json",
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert entry["directory_structure"] == {}
|
||||
|
||||
def test_file_name_map_parsed(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
fmap = {"hash1": "file1.txt"}
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"file_name_map": json.dumps(fmap),
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert entry["file_name_map"] == fmap
|
||||
|
||||
def test_faiss_missing_files_returns_no_file(
|
||||
self, internal_app, monkeypatch
|
||||
):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="faiss")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "no file"
|
||||
|
||||
def test_faiss_empty_filename_returns_no_file_name(
|
||||
self, internal_app, monkeypatch
|
||||
):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="faiss")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"file_faiss": (io.BytesIO(b""), ""),
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "no file name"
|
||||
|
||||
def test_remote_data_and_sync_frequency(self, internal_app, monkeypatch):
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "remote",
|
||||
"remote_data": '{"url":"http://example.com"}',
|
||||
"sync_frequency": "daily",
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert entry["sync_frequency"] == "daily"
|
||||
assert entry["remote_data"] == '{"url":"http://example.com"}'
|
||||
|
||||
def test_faiss_upload_with_valid_files(self, internal_app, monkeypatch):
|
||||
"""Cover lines 93-104: FAISS upload with both faiss and pkl files."""
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="faiss")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"file_faiss": (io.BytesIO(b"faiss data"), "index.faiss"),
|
||||
"file_pkl": (io.BytesIO(b"pkl data"), "index.pkl"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
mock_storage.save_file.assert_called()
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert entry is not None
|
||||
|
||||
def test_faiss_pkl_missing_returns_no_file(self, internal_app, monkeypatch):
|
||||
"""Cover lines 93-95: FAISS upload with faiss file but no pkl file."""
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="faiss")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"file_faiss": (io.BytesIO(b"faiss data"), "index.faiss"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "no file"
|
||||
|
||||
def test_faiss_pkl_empty_name_returns_no_file_name(self, internal_app, monkeypatch):
|
||||
"""Cover lines 97-98: FAISS upload with pkl but empty filename."""
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="faiss")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"file_faiss": (io.BytesIO(b"faiss data"), "index.faiss"),
|
||||
"file_pkl": (io.BytesIO(b""), ""),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "no file name"
|
||||
|
||||
def test_no_internal_key_rejects_upload(self, internal_app, monkeypatch):
|
||||
"""Verify that upload_index is rejected when INTERNAL_KEY is not set."""
|
||||
app, db = internal_app
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings",
|
||||
MagicMock(
|
||||
INTERNAL_KEY=None,
|
||||
UPLOAD_FOLDER="uploads",
|
||||
VECTOR_STORE="faiss",
|
||||
EMBEDDINGS_NAME="test",
|
||||
MONGO_DB_NAME="docsgpt",
|
||||
),
|
||||
)
|
||||
with app.test_client() as client:
|
||||
resp = client.post("/api/upload_index", data={"user": "attacker"})
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_update_existing_with_file_name_map(self, internal_app, monkeypatch):
|
||||
"""Cover line 124: update existing entry with file_name_map."""
|
||||
app, db = internal_app
|
||||
doc_id = ObjectId()
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
db["sources"].insert_one({"_id": doc_id, "user": "old_user", "name": "old"})
|
||||
|
||||
fmap = {"hash1": "file1.txt"}
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": str(doc_id),
|
||||
"type": "local",
|
||||
"file_name_map": json.dumps(fmap),
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": doc_id})
|
||||
assert entry["file_name_map"] == fmap
|
||||
|
||||
def test_invalid_file_name_map_defaults_none(self, internal_app, monkeypatch):
|
||||
"""Cover lines 77-79: invalid file_name_map JSON defaults to None."""
|
||||
app, db = internal_app
|
||||
doc_id = str(ObjectId())
|
||||
settings_mock = self._make_settings(vector_store="other")
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.settings", settings_mock
|
||||
)
|
||||
mock_storage = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"application.api.internal.routes.StorageCreator",
|
||||
MagicMock(get_storage=MagicMock(return_value=mock_storage)),
|
||||
)
|
||||
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/api/upload_index",
|
||||
data={
|
||||
"user": "u",
|
||||
"name": "n",
|
||||
"tokens": "0",
|
||||
"retriever": "classic",
|
||||
"id": doc_id,
|
||||
"type": "local",
|
||||
"file_name_map": "not valid json{{{",
|
||||
},
|
||||
headers=self._AUTH_HEADERS,
|
||||
)
|
||||
assert resp.json["status"] == "ok"
|
||||
|
||||
entry = db["sources"].find_one({"_id": ObjectId(doc_id)})
|
||||
assert "file_name_map" not in entry
|
||||
def test_faiss_saves_both_files(self, pg_conn):
|
||||
app = _make_app()
|
||||
fake_storage = MagicMock()
|
||||
data = {
|
||||
**self._base_form(source_id="legacy-src-7"),
|
||||
"file_faiss": (io.BytesIO(b"faiss-data"), "index.faiss"),
|
||||
"file_pkl": (io.BytesIO(b"pkl-data"), "index.pkl"),
|
||||
}
|
||||
with patch(
|
||||
"application.api.internal.routes.settings.INTERNAL_KEY", _TEST_KEY
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.VECTOR_STORE", "faiss"
|
||||
), patch(
|
||||
"application.api.internal.routes.settings.EMBEDDINGS_NAME", "emb"
|
||||
), patch(
|
||||
"application.api.internal.routes.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), _patch_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
r = c.post(
|
||||
"/api/upload_index",
|
||||
headers=_AUTH,
|
||||
data=data,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert r.json == {"status": "ok"}
|
||||
assert fake_storage.save_file.call_count == 2
|
||||
Reference in new issue
Block a user