mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 06:12:49 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
987 lines
36 KiB
Python
987 lines
36 KiB
Python
"""Tests for docsgpt/api/user/sources/upload.py."""
|
|
|
|
import io
|
|
import json
|
|
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):
|
|
@contextmanager
|
|
def _yield():
|
|
yield conn
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.db_session", _yield
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.db_readonly", _yield
|
|
):
|
|
yield
|
|
|
|
|
|
def _seed_source(pg_conn, user="u", name="src", **kw):
|
|
from docsgpt.storage.db.repositories.sources import SourcesRepository
|
|
return SourcesRepository(pg_conn).create(name, user_id=user, **kw)
|
|
|
|
|
|
class TestEnforceAudioPathSizeLimit:
|
|
def test_noop_for_non_audio(self, tmp_path):
|
|
from docsgpt.api.user.sources.upload import (
|
|
_enforce_audio_path_size_limit,
|
|
)
|
|
p = tmp_path / "doc.txt"
|
|
p.write_bytes(b"x" * 1024)
|
|
_enforce_audio_path_size_limit(str(p), "doc.txt")
|
|
|
|
def test_raises_for_large_audio(self, tmp_path):
|
|
from docsgpt.api.user.sources.upload import (
|
|
_enforce_audio_path_size_limit,
|
|
)
|
|
from docsgpt.stt.upload_limits import AudioFileTooLargeError
|
|
|
|
p = tmp_path / "audio.mp3"
|
|
p.write_bytes(b"x" * 100)
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.enforce_audio_file_size_limit",
|
|
side_effect=AudioFileTooLargeError("too large"),
|
|
):
|
|
with pytest.raises(AudioFileTooLargeError):
|
|
_enforce_audio_path_size_limit(str(p), "audio.mp3")
|
|
|
|
|
|
class TestUploadFile:
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
|
|
with app.test_request_context("/api/upload", method="POST"):
|
|
from flask import request
|
|
request.decoded_token = None
|
|
response = UploadFile().post()
|
|
assert response.status_code == 401
|
|
|
|
def test_returns_400_missing_fields(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
|
|
with app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={"user": "u"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_returns_400_empty_filenames(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
|
|
with app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "u", "name": "job",
|
|
"file": (io.BytesIO(b""), ""), # empty filename
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_uploads_single_file_successfully(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
|
|
fake_storage = MagicMock()
|
|
fake_task = MagicMock(id="task-1")
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "my_job",
|
|
"file": (io.BytesIO(b"content"), "doc.txt"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 200
|
|
assert response.json["success"] is True
|
|
assert response.json["task_id"] == "task-1"
|
|
|
|
def test_rejects_file_over_configured_upload_limit(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
|
|
fake_storage = MagicMock()
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.settings.UPLOAD_MAX_FILE_BYTES",
|
|
4,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
) as apply_async, app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "job",
|
|
"file": (io.BytesIO(b"12345"), "large.txt"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
|
|
assert response.status_code == 413
|
|
fake_storage.save_file.assert_not_called()
|
|
apply_async.assert_not_called()
|
|
|
|
def test_storage_error_returns_400(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
|
|
fake_storage = MagicMock()
|
|
fake_storage.save_file.side_effect = RuntimeError("boom")
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "j",
|
|
"file": (io.BytesIO(b"content"), "doc.txt"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_uploads_zip_extracts_files(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
import zipfile
|
|
|
|
# Build an in-memory zip containing 2 files
|
|
zip_buffer = io.BytesIO()
|
|
with zipfile.ZipFile(zip_buffer, "w") as zf:
|
|
zf.writestr("a.txt", "content a")
|
|
zf.writestr("sub/b.txt", "content b")
|
|
zip_buffer.seek(0)
|
|
|
|
fake_storage = MagicMock()
|
|
fake_task = MagicMock(id="task-zip")
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "job",
|
|
"file": (zip_buffer, "docs.zip"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 200
|
|
# save_file called at least twice (for 2 files in zip)
|
|
assert fake_storage.save_file.call_count >= 2
|
|
|
|
def test_rejects_zip_with_excessive_expansion_ratio(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
import zipfile
|
|
|
|
zip_buffer = io.BytesIO()
|
|
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
zf.writestr("bomb.txt", b"A" * (1024 * 1024))
|
|
zip_buffer.seek(0)
|
|
|
|
fake_storage = MagicMock()
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
) as apply_async, app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "job",
|
|
"file": (zip_buffer, "bomb.zip"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
|
|
assert response.status_code == 413
|
|
fake_storage.save_file.assert_not_called()
|
|
apply_async.assert_not_called()
|
|
|
|
def test_accepts_highly_compressible_csv(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
import zipfile
|
|
|
|
zip_buffer = io.BytesIO()
|
|
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as archive:
|
|
archive.writestr("export.csv", b"a,b,c\n" * 900_000)
|
|
zip_buffer.seek(0)
|
|
|
|
fake_storage = MagicMock()
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
return_value=MagicMock(id="csv-task"),
|
|
), app.test_request_context(
|
|
"/api/upload",
|
|
method="POST",
|
|
data={
|
|
"user": "alice",
|
|
"name": "job",
|
|
"file": (zip_buffer, "export.zip"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
|
|
assert response.status_code == 200
|
|
assert fake_storage.save_file.call_count == 1
|
|
|
|
def test_archive_error_names_file_and_escapes_log_controls(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
import zipfile
|
|
|
|
zip_buffer = io.BytesIO()
|
|
with zipfile.ZipFile(zip_buffer, "w") as archive:
|
|
archive.writestr("../bad\nFORGED", b"bad")
|
|
zip_buffer.seek(0)
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), patch.object(app.logger, "warning") as warning, app.test_request_context(
|
|
"/api/upload",
|
|
method="POST",
|
|
data={
|
|
"user": "alice",
|
|
"name": "job",
|
|
"file": (zip_buffer, "documents.zip"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
|
|
assert response.status_code == 413
|
|
assert response.json["message"] == (
|
|
"documents.zip: ZIP archive failed safety validation"
|
|
)
|
|
assert "bad" not in response.json["message"]
|
|
assert "\n" not in response.json["message"]
|
|
logged_values = " ".join(str(value) for value in warning.call_args.args)
|
|
assert "bad\\x0aFORGED" in logged_values
|
|
assert "\n" not in logged_values
|
|
|
|
def test_expands_nested_zip_before_storage(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
import zipfile
|
|
|
|
inner = io.BytesIO()
|
|
with zipfile.ZipFile(inner, "w") as archive:
|
|
archive.writestr("inside.txt", b"inside")
|
|
outer = io.BytesIO()
|
|
with zipfile.ZipFile(outer, "w") as archive:
|
|
archive.writestr("inner.zip", inner.getvalue())
|
|
outer.seek(0)
|
|
|
|
fake_storage = MagicMock()
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
return_value=MagicMock(id="nested-task"),
|
|
), app.test_request_context(
|
|
"/api/upload",
|
|
method="POST",
|
|
data={
|
|
"user": "alice",
|
|
"name": "job",
|
|
"file": (outer, "outer.zip"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
|
|
assert response.status_code == 200
|
|
saved_paths = [call.args[1] for call in fake_storage.save_file.call_args_list]
|
|
assert any(path.endswith("/inner/inside.txt") for path in saved_paths)
|
|
assert not any(path.endswith(".zip") for path in saved_paths)
|
|
|
|
def test_office_format_zip_saved_as_is(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
import zipfile
|
|
|
|
# .docx is technically a zip but should be saved as-is
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w") as zf:
|
|
zf.writestr("word/document.xml", "<x/>")
|
|
buf.seek(0)
|
|
|
|
fake_storage = MagicMock()
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload.ingest.apply_async",
|
|
return_value=MagicMock(id="t"),
|
|
), app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "job",
|
|
"file": (buf, "letter.docx"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 200
|
|
# Saved once (as the .docx directly, not extracted)
|
|
assert fake_storage.save_file.call_count == 1
|
|
|
|
def test_audio_too_large_returns_413(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadFile
|
|
from docsgpt.stt.upload_limits import AudioFileTooLargeError
|
|
|
|
fake_storage = MagicMock()
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.sources.upload._enforce_audio_path_size_limit",
|
|
side_effect=AudioFileTooLargeError("too large"),
|
|
), app.test_request_context(
|
|
"/api/upload", method="POST",
|
|
data={
|
|
"user": "alice", "name": "j",
|
|
"file": (io.BytesIO(b"audio-content"), "song.mp3"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "alice"}
|
|
response = UploadFile().post()
|
|
assert response.status_code == 413
|
|
|
|
|
|
class TestUploadRemote:
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadRemote
|
|
|
|
with app.test_request_context(
|
|
"/api/remote", method="POST",
|
|
data={"user": "u", "source": "github"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = None
|
|
response = UploadRemote().post()
|
|
assert response.status_code == 401
|
|
|
|
def test_returns_missing_fields(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadRemote
|
|
|
|
with app.test_request_context(
|
|
"/api/remote", method="POST",
|
|
data={"user": "u"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadRemote().post()
|
|
# check_required_fields returns a response; status is 400
|
|
# The response is returned directly by missing_fields branch
|
|
assert response.status_code == 400
|
|
|
|
def test_uploads_github_remote_success(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadRemote
|
|
|
|
fake_task = MagicMock(id="remote-task-1")
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.ingest_remote.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/remote", method="POST",
|
|
data={
|
|
"user": "u", "source": "github", "name": "gh",
|
|
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadRemote().post()
|
|
assert response.status_code == 200
|
|
assert response.json["task_id"] == "remote-task-1"
|
|
|
|
def test_uploads_url_source(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadRemote
|
|
|
|
fake_task = MagicMock(id="url-task")
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.ingest_remote.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/remote", method="POST",
|
|
data={
|
|
"user": "u", "source": "crawler", "name": "crawl",
|
|
"data": json.dumps({"url": "https://example.com"}),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadRemote().post()
|
|
assert response.status_code == 200
|
|
|
|
def test_uploads_reddit_source(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadRemote
|
|
|
|
fake_task = MagicMock(id="reddit-task")
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.ingest_remote.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/remote", method="POST",
|
|
data={
|
|
"user": "u", "source": "reddit", "name": "r",
|
|
"data": json.dumps({"subreddit": "python"}),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadRemote().post()
|
|
assert response.status_code == 200
|
|
|
|
def test_upload_exception_returns_400(self, app):
|
|
from docsgpt.api.user.sources.upload import UploadRemote
|
|
|
|
with patch(
|
|
"docsgpt.api.user.sources.upload.ingest_remote.apply_async",
|
|
side_effect=RuntimeError("boom"),
|
|
), app.test_request_context(
|
|
"/api/remote", method="POST",
|
|
data={
|
|
"user": "u", "source": "github", "name": "x",
|
|
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = UploadRemote().post()
|
|
assert response.status_code == 400
|
|
|
|
|
|
class TestManageSourceFiles:
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
with app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={"source_id": "x", "operation": "add"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = None
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 401
|
|
|
|
def test_returns_400_missing_required(self, app):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
with app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={"source_id": "x"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_returns_400_invalid_operation(self, app):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
with app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={"source_id": "x", "operation": "weird"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_returns_404_source_not_found(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": "00000000-0000-0000-0000-000000000000",
|
|
"operation": "add",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 404
|
|
|
|
def test_rejects_bad_parent_dir(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-bad-parent"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "add",
|
|
"parent_dir": "/abs-path",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_add_no_files_returns_400(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-add-nofile"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={"source_id": str(src["id"]), "operation": "add"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_add_rejects_file_over_upload_limit(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-add-large"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data/src")
|
|
fake_storage = MagicMock()
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.upload_limits.settings.UPLOAD_MAX_FILE_BYTES", 4
|
|
), patch(
|
|
"docsgpt.api.user.tasks.reingest_source_task.apply_async"
|
|
) as apply_async, app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "add",
|
|
"file": (io.BytesIO(b"12345"), "large.txt"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
|
|
assert response.status_code == 413
|
|
fake_storage.save_file.assert_not_called()
|
|
apply_async.assert_not_called()
|
|
|
|
def test_add_files_success(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-add-ok"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data/src")
|
|
|
|
fake_storage = MagicMock()
|
|
fake_task = MagicMock(id="reingest-1")
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.tasks.reingest_source_task.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "add",
|
|
"file": (io.BytesIO(b"content"), "new.txt"),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 200
|
|
assert response.json["success"] is True
|
|
assert "new.txt" in response.json["added_files"]
|
|
|
|
def test_remove_missing_file_paths_returns_400(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rm-nolist"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={"source_id": str(src["id"]), "operation": "remove"},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_remove_invalid_json_file_paths(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rm-bad"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove",
|
|
"file_paths": "not-json",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_remove_rejects_path_traversal(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rm-trav"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove",
|
|
"file_paths": json.dumps(["../escape.txt"]),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_remove_files_success(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rm-ok"
|
|
src = _seed_source(
|
|
pg_conn, user=user,
|
|
file_path="/data/src",
|
|
file_name_map={"a.txt": "Original A.txt"},
|
|
)
|
|
|
|
fake_storage = MagicMock()
|
|
fake_storage.file_exists.return_value = True
|
|
fake_task = MagicMock(id="reingest-rm")
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.tasks.reingest_source_task.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove",
|
|
"file_paths": json.dumps(["a.txt"]),
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 200
|
|
assert "a.txt" in response.json["removed_files"]
|
|
|
|
def test_remove_directory_missing_path(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rmdir-missing"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove_directory",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_remove_directory_rejects_bad_path(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rmdir-bad"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=MagicMock(),
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove_directory",
|
|
"directory_path": "../escape",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_remove_directory_404_when_not_directory(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rmdir-notdir"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
fake_storage = MagicMock()
|
|
fake_storage.is_directory.return_value = False
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove_directory",
|
|
"directory_path": "subdir",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 404
|
|
|
|
def test_remove_directory_success(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rmdir-ok"
|
|
src = _seed_source(
|
|
pg_conn, user=user, file_path="/data",
|
|
file_name_map={"sub/a.txt": "A"},
|
|
)
|
|
|
|
fake_storage = MagicMock()
|
|
fake_storage.is_directory.return_value = True
|
|
fake_storage.remove_directory.return_value = True
|
|
fake_task = MagicMock(id="reingest-dir")
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), patch(
|
|
"docsgpt.api.user.tasks.reingest_source_task.apply_async",
|
|
return_value=fake_task,
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove_directory",
|
|
"directory_path": "sub",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 200
|
|
assert response.json["removed_directory"] == "sub"
|
|
|
|
def test_remove_directory_storage_failure_returns_500(self, app, pg_conn):
|
|
from docsgpt.api.user.sources.upload import ManageSourceFiles
|
|
|
|
user = "u-rmdir-fail"
|
|
src = _seed_source(pg_conn, user=user, file_path="/data")
|
|
|
|
fake_storage = MagicMock()
|
|
fake_storage.is_directory.return_value = True
|
|
fake_storage.remove_directory.return_value = False
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.sources.upload.StorageCreator.get_storage",
|
|
return_value=fake_storage,
|
|
), app.test_request_context(
|
|
"/api/manage_source_files", method="POST",
|
|
data={
|
|
"source_id": str(src["id"]),
|
|
"operation": "remove_directory",
|
|
"directory_path": "sub",
|
|
},
|
|
content_type="multipart/form-data",
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = ManageSourceFiles().post()
|
|
assert response.status_code == 500
|
|
|
|
|
|
class TestTaskStatus:
|
|
def test_returns_400_missing_task_id(self, app):
|
|
from docsgpt.api.user.sources.upload import TaskStatus
|
|
|
|
with app.test_request_context("/api/task_status"):
|
|
response = TaskStatus().get()
|
|
assert response.status_code == 400
|
|
|
|
def test_returns_task_status(self, app):
|
|
from docsgpt.api.user.sources.upload import TaskStatus
|
|
|
|
fake_task = MagicMock()
|
|
fake_task.status = "SUCCESS"
|
|
fake_task.info = {"result": "ok"}
|
|
|
|
fake_celery = MagicMock()
|
|
fake_celery.AsyncResult.return_value = fake_task
|
|
|
|
with patch(
|
|
"docsgpt.celery_init.celery", fake_celery
|
|
), app.test_request_context("/api/task_status?task_id=t-123"):
|
|
response = TaskStatus().get()
|
|
assert response.status_code == 200
|
|
assert response.json["status"] == "SUCCESS"
|
|
|
|
def test_pending_without_workers_returns_503(self, app):
|
|
from docsgpt.api.user.sources.upload import TaskStatus
|
|
|
|
fake_task = MagicMock()
|
|
fake_task.status = "PENDING"
|
|
fake_task.info = None
|
|
|
|
fake_inspect = MagicMock()
|
|
fake_inspect.ping.return_value = None # no workers
|
|
|
|
fake_celery = MagicMock()
|
|
fake_celery.AsyncResult.return_value = fake_task
|
|
fake_celery.control.inspect.return_value = fake_inspect
|
|
|
|
with patch(
|
|
"docsgpt.celery_init.celery", fake_celery
|
|
), app.test_request_context("/api/task_status?task_id=t-999"):
|
|
response = TaskStatus().get()
|
|
assert response.status_code == 503
|
|
|
|
def test_exception_returns_400(self, app):
|
|
from docsgpt.api.user.sources.upload import TaskStatus
|
|
|
|
fake_celery = MagicMock()
|
|
fake_celery.AsyncResult.side_effect = RuntimeError("boom")
|
|
|
|
with patch(
|
|
"docsgpt.celery_init.celery", fake_celery
|
|
), app.test_request_context("/api/task_status?task_id=t-err"):
|
|
response = TaskStatus().get()
|
|
assert response.status_code == 400
|
|
|
|
def test_non_serializable_info_gets_stringified(self, app):
|
|
from docsgpt.api.user.sources.upload import TaskStatus
|
|
|
|
class WeirdObj:
|
|
def __str__(self):
|
|
return "weird-str"
|
|
|
|
fake_task = MagicMock()
|
|
fake_task.status = "SUCCESS"
|
|
fake_task.info = WeirdObj()
|
|
|
|
fake_celery = MagicMock()
|
|
fake_celery.AsyncResult.return_value = fake_task
|
|
|
|
with patch(
|
|
"docsgpt.celery_init.celery", fake_celery
|
|
), app.test_request_context("/api/task_status?task_id=t-weird"):
|
|
response = TaskStatus().get()
|
|
assert response.status_code == 200
|
|
assert response.json["result"] == "weird-str"
|