mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 20:13:04 +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.
457 lines
16 KiB
Python
457 lines
16 KiB
Python
import io
|
|
import os
|
|
from unittest.mock import MagicMock, mock_open, patch
|
|
|
|
import pytest
|
|
from docsgpt.storage.local import LocalStorage
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_base_dir():
|
|
return os.path.realpath("/tmp/test_storage")
|
|
|
|
|
|
@pytest.fixture
|
|
def local_storage(temp_base_dir):
|
|
return LocalStorage(base_dir=temp_base_dir)
|
|
|
|
|
|
@pytest.fixture
|
|
def real_storage(tmp_path):
|
|
"""Storage over a real directory, for the write paths worth exercising."""
|
|
base = os.path.realpath(str(tmp_path))
|
|
return LocalStorage(base_dir=base), base
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageInitialization:
|
|
|
|
def test_init_with_custom_base_dir(self):
|
|
storage = LocalStorage(base_dir="/custom/path")
|
|
assert storage.base_dir == "/custom/path"
|
|
|
|
def test_init_with_default_base_dir(self):
|
|
storage = LocalStorage()
|
|
assert storage.base_dir is not None
|
|
assert isinstance(storage.base_dir, str)
|
|
|
|
def test_get_full_path_with_relative_path(self, local_storage):
|
|
result = local_storage._get_full_path("documents/test.txt")
|
|
expected = os.path.realpath(os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt"))
|
|
assert result == expected
|
|
|
|
def test_get_full_path_with_absolute_path_outside_base_raises(self, local_storage):
|
|
with pytest.raises(ValueError, match="Path traversal detected"):
|
|
local_storage._get_full_path("/absolute/path/test.txt")
|
|
|
|
def test_save_file_creates_directory_and_saves(self, real_storage):
|
|
storage, base = real_storage
|
|
result = storage.save_file(io.BytesIO(b"test content"), "documents/test.txt")
|
|
|
|
written = os.path.join(base, "documents/test.txt")
|
|
assert os.path.isfile(written)
|
|
with open(written, "rb") as f:
|
|
assert f.read() == b"test content"
|
|
assert result == {"storage_type": "local"}
|
|
|
|
def test_save_file_with_save_method(self, real_storage):
|
|
"""Werkzeug's ``FileStorage.save`` accepts the destination handle."""
|
|
storage, base = real_storage
|
|
|
|
class _Uploaded:
|
|
def save(self, dst):
|
|
dst.write(b"from save()")
|
|
|
|
result = storage.save_file(_Uploaded(), "documents/test.txt")
|
|
|
|
with open(os.path.join(base, "documents/test.txt"), "rb") as f:
|
|
assert f.read() == b"from save()"
|
|
assert result == {"storage_type": "local"}
|
|
|
|
def test_save_file_replaces_existing_content(self, real_storage):
|
|
storage, base = real_storage
|
|
storage.save_file(io.BytesIO(b"first"), "a/b.bin")
|
|
storage.save_file(io.BytesIO(b"second"), "a/b.bin")
|
|
|
|
with open(os.path.join(base, "a/b.bin"), "rb") as f:
|
|
assert f.read() == b"second"
|
|
|
|
def test_failed_write_leaves_the_previous_file_intact(self, real_storage):
|
|
"""The whole point of writing through a temp file.
|
|
|
|
``reembed`` rewrites an index in place; a half-written ``index.faiss``
|
|
loads at neither the old width nor the new one, so an interrupted write
|
|
must leave the previous bytes untouched rather than truncate them.
|
|
"""
|
|
storage, base = real_storage
|
|
storage.save_file(io.BytesIO(b"the good index"), "indexes/s1/index.faiss")
|
|
|
|
class _DiesHalfway:
|
|
def __init__(self):
|
|
self._served = False
|
|
|
|
def read(self, size=-1):
|
|
if self._served:
|
|
raise OSError("connection reset")
|
|
self._served = True
|
|
return b"corrupt"
|
|
|
|
with pytest.raises(OSError, match="connection reset"):
|
|
storage.save_file(_DiesHalfway(), "indexes/s1/index.faiss")
|
|
|
|
with open(os.path.join(base, "indexes/s1/index.faiss"), "rb") as f:
|
|
assert f.read() == b"the good index"
|
|
|
|
def test_failed_write_leaves_no_temp_file_behind(self, real_storage):
|
|
storage, base = real_storage
|
|
|
|
class _Explodes:
|
|
def read(self, size=-1):
|
|
raise OSError("boom")
|
|
|
|
with pytest.raises(OSError):
|
|
storage.save_file(_Explodes(), "indexes/s1/index.faiss")
|
|
|
|
assert os.listdir(os.path.join(base, "indexes/s1")) == []
|
|
|
|
def test_save_file_keeps_the_existing_permissions(self, real_storage):
|
|
"""``mkstemp`` is 0600; the replacement must not silently narrow access."""
|
|
storage, base = real_storage
|
|
storage.save_file(io.BytesIO(b"one"), "a/b.bin")
|
|
target = os.path.join(base, "a/b.bin")
|
|
os.chmod(target, 0o640)
|
|
|
|
storage.save_file(io.BytesIO(b"two"), "a/b.bin")
|
|
|
|
assert os.stat(target).st_mode & 0o777 == 0o640
|
|
|
|
def test_save_file_with_absolute_path_outside_base_raises(self, local_storage):
|
|
file_data = io.BytesIO(b"test content")
|
|
path = "/absolute/path/test.txt"
|
|
|
|
with pytest.raises(ValueError, match="Path traversal detected"):
|
|
local_storage.save_file(file_data, path)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageGetFile:
|
|
|
|
@patch("os.path.exists", return_value=True)
|
|
@patch("builtins.open", new_callable=mock_open, read_data=b"file content")
|
|
def test_get_file_returns_file_handle(self, mock_file, mock_exists, local_storage):
|
|
path = "documents/test.txt"
|
|
|
|
result = local_storage.get_file(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert mock_file.call_count == 1
|
|
assert os.path.normpath(mock_file.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert result is not None
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
def test_get_file_raises_error_when_not_found(self, mock_exists, local_storage):
|
|
path = "documents/nonexistent.txt"
|
|
|
|
with pytest.raises(FileNotFoundError, match="File not found"):
|
|
local_storage.get_file(path)
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/nonexistent.txt")
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageFileSize:
|
|
|
|
@patch("os.path.getsize", return_value=42)
|
|
def test_get_file_size_uses_metadata_without_opening(
|
|
self, mock_getsize, local_storage
|
|
):
|
|
assert local_storage.get_file_size("avatars/test.png") == 42
|
|
|
|
expected_path = os.path.join(
|
|
os.path.realpath("/tmp/test_storage"), "avatars/test.png"
|
|
)
|
|
assert os.path.normpath(mock_getsize.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageDeleteFile:
|
|
|
|
@patch("os.remove")
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_delete_file_removes_existing_file(
|
|
self, mock_exists, mock_remove, local_storage
|
|
):
|
|
path = "documents/test.txt"
|
|
|
|
result = local_storage.delete_file(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
|
|
assert result is True
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert mock_remove.call_count == 1
|
|
assert os.path.normpath(mock_remove.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
def test_delete_file_returns_false_when_not_found(self, mock_exists, local_storage):
|
|
path = "documents/nonexistent.txt"
|
|
|
|
result = local_storage.delete_file(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/nonexistent.txt")
|
|
assert result is False
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageFileExists:
|
|
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_file_exists_returns_true_when_file_found(self, mock_exists, local_storage):
|
|
path = "documents/test.txt"
|
|
|
|
result = local_storage.file_exists(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
|
|
assert result is True
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
def test_file_exists_returns_false_when_not_found(self, mock_exists, local_storage):
|
|
path = "documents/nonexistent.txt"
|
|
|
|
result = local_storage.file_exists(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/nonexistent.txt")
|
|
assert result is False
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageListFiles:
|
|
|
|
@patch("os.walk")
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_list_files_returns_all_files_in_directory(
|
|
self, mock_exists, mock_walk, local_storage
|
|
):
|
|
directory = "documents"
|
|
base_dir = os.path.join(os.path.realpath("/tmp/test_storage"), "documents")
|
|
|
|
mock_walk.return_value = [
|
|
(base_dir, ["subdir"], ["file1.txt", "file2.txt"]),
|
|
(os.path.join(base_dir, "subdir"), [], ["file3.txt"]),
|
|
]
|
|
|
|
result = local_storage.list_files(directory)
|
|
|
|
assert len(result) == 3
|
|
result_normalized = [os.path.normpath(f) for f in result]
|
|
assert os.path.normpath("documents/file1.txt") in result_normalized
|
|
assert os.path.normpath("documents/file2.txt") in result_normalized
|
|
assert os.path.normpath("documents/subdir/file3.txt") in result_normalized
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
def test_list_files_returns_empty_list_when_directory_not_found(
|
|
self, mock_exists, local_storage
|
|
):
|
|
directory = "nonexistent"
|
|
|
|
result = local_storage.list_files(directory)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "nonexistent")
|
|
assert result == []
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageProcessFile:
|
|
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_process_file_calls_processor_with_full_path(
|
|
self, mock_exists, local_storage
|
|
):
|
|
path = "documents/test.txt"
|
|
processor_func = MagicMock(return_value="processed")
|
|
|
|
result = local_storage.process_file(path, processor_func, extra_arg="value")
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
|
|
assert result == "processed"
|
|
assert processor_func.call_count == 1
|
|
call_kwargs = processor_func.call_args[1]
|
|
assert os.path.normpath(call_kwargs["local_path"]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert call_kwargs["extra_arg"] == "value"
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
def test_process_file_raises_error_when_file_not_found(
|
|
self, mock_exists, local_storage
|
|
):
|
|
path = "documents/nonexistent.txt"
|
|
processor_func = MagicMock()
|
|
|
|
with pytest.raises(FileNotFoundError, match="File not found"):
|
|
local_storage.process_file(path, processor_func)
|
|
processor_func.assert_not_called()
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageIsDirectory:
|
|
|
|
@patch("os.path.isdir", return_value=True)
|
|
def test_is_directory_returns_true_when_directory_exists(
|
|
self, mock_isdir, local_storage
|
|
):
|
|
path = "documents"
|
|
|
|
result = local_storage.is_directory(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents")
|
|
assert result is True
|
|
assert mock_isdir.call_count == 1
|
|
assert os.path.normpath(mock_isdir.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("os.path.isdir", return_value=False)
|
|
def test_is_directory_returns_false_when_not_directory(
|
|
self, mock_isdir, local_storage
|
|
):
|
|
path = "documents/test.txt"
|
|
|
|
result = local_storage.is_directory(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
|
|
assert result is False
|
|
assert mock_isdir.call_count == 1
|
|
assert os.path.normpath(mock_isdir.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLocalStorageRemoveDirectory:
|
|
|
|
@patch("shutil.rmtree")
|
|
@patch("os.path.isdir", return_value=True)
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_remove_directory_deletes_directory(
|
|
self, mock_exists, mock_isdir, mock_rmtree, local_storage
|
|
):
|
|
directory = "documents"
|
|
|
|
result = local_storage.remove_directory(directory)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents")
|
|
assert result is True
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert mock_isdir.call_count == 1
|
|
assert os.path.normpath(mock_isdir.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert mock_rmtree.call_count == 1
|
|
assert os.path.normpath(mock_rmtree.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
def test_remove_directory_returns_false_when_not_exists(
|
|
self, mock_exists, local_storage
|
|
):
|
|
directory = "nonexistent"
|
|
|
|
result = local_storage.remove_directory(directory)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "nonexistent")
|
|
assert result is False
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("os.path.isdir", return_value=False)
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_remove_directory_returns_false_when_not_directory(
|
|
self, mock_exists, mock_isdir, local_storage
|
|
):
|
|
path = "documents/test.txt"
|
|
|
|
result = local_storage.remove_directory(path)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
|
|
assert result is False
|
|
assert mock_exists.call_count == 1
|
|
assert os.path.normpath(mock_exists.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
assert mock_isdir.call_count == 1
|
|
assert os.path.normpath(mock_isdir.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("shutil.rmtree", side_effect=OSError("Permission denied"))
|
|
@patch("os.path.isdir", return_value=True)
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_remove_directory_returns_false_on_os_error(
|
|
self, mock_exists, mock_isdir, mock_rmtree, local_storage
|
|
):
|
|
directory = "documents"
|
|
|
|
result = local_storage.remove_directory(directory)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents")
|
|
assert result is False
|
|
assert mock_rmtree.call_count == 1
|
|
assert os.path.normpath(mock_rmtree.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|
|
|
|
@patch("shutil.rmtree", side_effect=PermissionError("Access denied"))
|
|
@patch("os.path.isdir", return_value=True)
|
|
@patch("os.path.exists", return_value=True)
|
|
def test_remove_directory_returns_false_on_permission_error(
|
|
self, mock_exists, mock_isdir, mock_rmtree, local_storage
|
|
):
|
|
directory = "documents"
|
|
|
|
result = local_storage.remove_directory(directory)
|
|
|
|
expected_path = os.path.join(os.path.realpath("/tmp/test_storage"), "documents")
|
|
assert result is False
|
|
assert mock_rmtree.call_count == 1
|
|
assert os.path.normpath(mock_rmtree.call_args[0][0]) == os.path.normpath(
|
|
expected_path
|
|
)
|