diff --git a/backend/open_webui/test/apps/webui/storage/test_provider.py b/backend/open_webui/test/apps/webui/storage/test_provider.py index 8df07a9fb..8d1090541 100644 --- a/backend/open_webui/test/apps/webui/storage/test_provider.py +++ b/backend/open_webui/test/apps/webui/storage/test_provider.py @@ -1,7 +1,17 @@ +import io + import pytest from open_webui.storage import provider +def mock_upload_dir(monkeypatch, tmp_path): + """Fixture to monkey-patch the UPLOAD_DIR and create a temporary directory.""" + directory = tmp_path / "uploads" + directory.mkdir() + monkeypatch.setattr(provider, "UPLOAD_DIR", str(directory)) + return directory + + def test_imports(): provider.StorageProvider provider.LocalStorageProvider @@ -17,36 +27,83 @@ def test_get_storage_provider(): with pytest.raises(RuntimeError): provider.get_storage_provider("invalid") + def test_class_instantiation(): with pytest.raises(TypeError): provider.StorageProvider() with pytest.raises(TypeError): + class Test(provider.StorageProvider): pass + Test() provider.LocalStorageProvider() provider.S3StorageProvider() class TestLocalStorageProvider(provider.LocalStorageProvider): - def test_upload_file(self): - pass + file_content = b"test content" + file_bytesio = io.BytesIO(file_content) + filename = "test.txt" + filename_extra = "test_exyta.txt" + file_bytesio_empty = io.BytesIO() + + def test_upload_file(self, monkeypatch, tmp_path): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + contents, file_path = self.upload_file(self.file_bytesio, self.filename) + assert (upload_dir / self.filename).exists() + assert (upload_dir / self.filename).read_bytes() == self.file_content + assert contents == self.file_content + assert file_path == str(upload_dir / self.filename) + with pytest.raises(ValueError): + self.upload_file(self.file_bytesio_empty, self.filename) + + def test_get_file(self, monkeypatch, tmp_path): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + file_path = str(upload_dir / self.filename) + file_path_return = self.get_file(file_path) + assert file_path == file_path_return + + def test_delete_file(self, monkeypatch, tmp_path): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + (upload_dir / self.filename).write_bytes(self.file_content) + assert (upload_dir / self.filename).exists() + file_path = str(upload_dir / self.filename) + self.delete_file(file_path) + assert not (upload_dir / self.filename).exists() + + def test_delete_all_files(self, monkeypatch, tmp_path): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + (upload_dir / self.filename).write_bytes(self.file_content) + (upload_dir / self.filename_extra).write_bytes(self.file_content) + self.delete_all_files() + assert not (upload_dir / self.filename).exists() + assert not (upload_dir / self.filename_extra).exists() + + +class TestS3StorageProvider(provider.S3StorageProvider): + file_content = b"test content" + file_bytesio = io.BytesIO(file_content) + filename = "test.txt" + filename_extra = "test_extra.txt" + file_bytesio_empty = io.BytesIO() + bucket_name = "my-bucket" + + def test_upload_file(self, monkeypatch, tmp_path): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + contents, file_path = self.upload_file(self.file_bytesio, self.filename) + assert (upload_dir / self.filename).exists() + assert (upload_dir / self.filename).read_bytes() == self.file_content + assert contents == self.file_content + assert file_path == str(upload_dir / self.filename) + with pytest.raises(ValueError): + self.upload_file(self.file_bytesio_empty, self.filename) + def test_get_file(self): pass + def test_delete_file(self): pass + def test_delete_all_files(self): pass - - -class TestLocalStorageProvider(provider.S3StorageProvider): - def test_upload_file(self): - pass - def test_get_file(self): - pass - def test_delete_file(self): - pass - def test_delete_all_files(self): - pass - -)