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 2bc3ccb64..08338375e 100644 --- a/backend/open_webui/test/apps/webui/storage/test_provider.py +++ b/backend/open_webui/test/apps/webui/storage/test_provider.py @@ -282,26 +282,8 @@ class TestGCSStorageProvider: class TestAzureStorageProvider: - def __init__(self): - super().__init__() - - @pytest.fixture(scope="class") - def setup(self, monkeypatch): - """Mock BlobServiceClient and BlobContainerClient for local testing""" - # Create mock Blob Service Client - mock_blob_service_client = MagicMock() - mock_container_client = MagicMock() - mock_blob_client = MagicMock() - - # Set up return values - mock_blob_service_client.get_container_client.return_value = mock_container_client - mock_container_client.get_blob_client.return_value = mock_blob_client - - # Mock `from_connection_string` and `BlobServiceClient` constructor - monkeypatch.setattr(azure.storage.blob.BlobServiceClient, lambda *_: mock_blob_service_client) - monkeypatch.setattr(azure.storage.blob.BlobContainerClient, lambda *_: mock_container_client) - monkeypatch.setattr(azure.storage.blob.BlobClient, lambda *_: mock_blob_client) - + @pytest.fixture(autouse=True) + def setup_storage(self, monkeypatch): self.Storage = provider.AzureStorageProvider() self.Storage.endpoint = "https://myaccount.blob.core.windows.net" self.Storage.container_name = "my-container" @@ -310,34 +292,48 @@ class TestAzureStorageProvider: self.filename_extra = "test_extra.txt" self.file_bytesio_empty = io.BytesIO() - # Apply to instance variables + # Create mock Blob Service Client and related clients + mock_blob_service_client = MagicMock() + mock_container_client = MagicMock() + mock_blob_client = MagicMock() + + # Set up return values for the mock + mock_blob_service_client.get_container_client.return_value = mock_container_client + mock_container_client.get_blob_client.return_value = mock_blob_client + + # Monkeypatch the Azure classes to return our mocks + monkeypatch.setattr( + azure.storage.blob, "BlobServiceClient", lambda *args, **kwargs: mock_blob_service_client + ) + monkeypatch.setattr( + azure.storage.blob, "BlobContainerClient", lambda *args, **kwargs: mock_container_client + ) + monkeypatch.setattr( + azure.storage.blob, "BlobClient", lambda *args, **kwargs: mock_blob_client + ) + + # Apply mocks to the Storage instance self.Storage.blob_service_client = mock_blob_service_client self.Storage.container_client = mock_container_client - yield - - def test_upload_file(self, monkeypatch, tmp_path, setup): - """Test uploading a file to mocked Azure Storage.""" + def test_upload_file(self, monkeypatch, tmp_path): upload_dir = mock_upload_dir(monkeypatch, tmp_path) - + # Simulate an error when container does not exist self.Storage.container_client.get_blob_client.side_effect = Exception("Container does not exist") - with pytest.raises(Exception): self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) # Reset side effect and create container self.Storage.container_client.get_blob_client.side_effect = None - self.Storage.create_container() - contents, azure_file_path = self.Storage.upload_file( - io.BytesIO(self.file_content), self.filename - ) + contents, azure_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) # Assertions self.Storage.container_client.get_blob_client.assert_called_with(self.filename) - self.Storage.container_client.get_blob_client().upload_blob.assert_called_once_with(self.file_content, overwrite=True) - + self.Storage.container_client.get_blob_client().upload_blob.assert_called_once_with( + self.file_content, overwrite=True + ) assert contents == self.file_content assert azure_file_path == f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}" assert (upload_dir / self.filename).exists() @@ -346,42 +342,38 @@ class TestAzureStorageProvider: with pytest.raises(ValueError): self.Storage.upload_file(self.file_bytesio_empty, self.filename) - def test_get_file(self, monkeypatch, tmp_path, setup): - """Test retrieving a file from mocked Azure Storage.""" + def test_get_file(self, monkeypatch, tmp_path): upload_dir = mock_upload_dir(monkeypatch, tmp_path) self.Storage.create_container() # Mock upload behavior self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) - # Mock blob download behavior self.Storage.container_client.get_blob_client().download_blob().readall.return_value = self.file_content - file_path = self.Storage.get_file(f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}") + file_url = f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}" + file_path = self.Storage.get_file(file_url) assert file_path == str(upload_dir / self.filename) assert (upload_dir / self.filename).exists() assert (upload_dir / self.filename).read_bytes() == self.file_content - def test_delete_file(self, monkeypatch, tmp_path, setup): - """Test deleting a file from mocked Azure Storage.""" + def test_delete_file(self, monkeypatch, tmp_path): upload_dir = mock_upload_dir(monkeypatch, tmp_path) self.Storage.create_container() - # Mock upload + # Mock file upload self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) - # Mock deletion self.Storage.container_client.get_blob_client().delete_blob.return_value = None - self.Storage.delete_file(f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}") + file_url = f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}" + self.Storage.delete_file(file_url) - # Assertions self.Storage.container_client.get_blob_client().delete_blob.assert_called_once() assert not (upload_dir / self.filename).exists() - def test_delete_all_files(self, monkeypatch, tmp_path, setup): - """Test deleting all files from mocked Azure Storage.""" + def test_delete_all_files(self, monkeypatch, tmp_path): upload_dir = mock_upload_dir(monkeypatch, tmp_path) self.Storage.create_container() @@ -396,22 +388,18 @@ class TestAzureStorageProvider: ] self.Storage.container_client.get_blob_client().delete_blob.return_value = None - # Call delete all files self.Storage.delete_all_files() - # Assertions self.Storage.container_client.list_blobs.assert_called_once() self.Storage.container_client.get_blob_client().delete_blob.assert_any_call() assert not (upload_dir / self.filename).exists() assert not (upload_dir / self.filename_extra).exists() - def test_get_file_not_found(self, monkeypatch, setup): - """Test handling when a requested file does not exist.""" + def test_get_file_not_found(self, monkeypatch): self.Storage.create_container() - # Mock behavior to raise an error for missing files + file_url = f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}" + # Mock behavior to raise an error for missing blobs self.Storage.container_client.get_blob_client().download_blob.side_effect = Exception("Blob not found") - with pytest.raises(Exception, match="Blob not found"): - self.Storage.get_file(f"https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}") - + self.Storage.get_file(file_url) \ No newline at end of file