Update test_provider.py

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)