diff --git a/.github/workflows/integration-test.disabled b/.github/workflows/integration-test.disabled index b248df4b5..946735c5b 100644 --- a/.github/workflows/integration-test.disabled +++ b/.github/workflows/integration-test.disabled @@ -5,6 +5,7 @@ on: branches: - main - dev + - azure-storage pull_request: branches: - main 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 4c3112526..5cfb1b6a8 100644 --- a/backend/open_webui/test/apps/webui/storage/test_provider.py +++ b/backend/open_webui/test/apps/webui/storage/test_provider.py @@ -3,10 +3,11 @@ import os import boto3 import pytest from botocore.exceptions import ClientError -from moto import mock_aws +from moto import mock_aws, mock_azure from open_webui.storage import provider from gcp_storage_emulator.server import create_server from google.cloud import storage +from azure.storage.blob import BlobServiceClient def mock_upload_dir(monkeypatch, tmp_path): @@ -276,3 +277,100 @@ class TestGCSStorageProvider: assert not (upload_dir / self.filename_extra).exists() assert self.Storage.bucket.get_blob(self.filename) == None assert self.Storage.bucket.get_blob(self.filename_extra) == None + + +class TestAzureStorageProvider: + def __init__(self): + self.Storage = provider.AzureStorageProvider() + self.Storage.container_name = "my-container" + self.file_content = b"test content" + self.filename = "test.txt" + self.filename_extra = "test_exyta.txt" + self.file_bytesio_empty = io.BytesIO() + super().__init__() + + @pytest.fixture(scope="class") + def setup(self, monkeypatch): + connection_string = "DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtl6rE4rWlgEoMF1rA==;BlobEndpoint=http://127.0.0.1:10000/devstoreaccount1;" + self.Storage.blob_service_client = BlobServiceClient.from_connection_string(connection_string) + self.Storage.container_client = self.Storage.blob_service_client.get_container_client(self.Storage.container_name) + monkeypatch.setattr(self.Storage, "blob_service_client", self.Storage.blob_service_client) + monkeypatch.setattr(self.Storage, "container_client", self.Storage.container_client) + yield + self.Storage.container_client.delete_container() + + def test_upload_file(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + # Azure checks + with pytest.raises(Exception): + self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) + self.Storage.create_container() + contents, azure_file_path = self.Storage.upload_file( + io.BytesIO(self.file_content), self.filename + ) + blob = self.Storage.blob_service_client.get_blob_client( + container=self.Storage.container_name, blob=self.filename + ) + assert self.file_content == blob.download_blob().readall() + # local checks + assert (upload_dir / self.filename).exists() + assert (upload_dir / self.filename).read_bytes() == self.file_content + assert contents == self.file_content + assert azure_file_path == "azure://" + self.Storage.container_name + "/" + self.filename + with pytest.raises(ValueError): + self.Storage.upload_file(self.file_bytesio_empty, self.filename) + + def test_get_file(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + self.Storage.create_container() + contents, azure_file_path = self.Storage.upload_file( + io.BytesIO(self.file_content), self.filename + ) + file_path = self.Storage.get_file(azure_file_path) + assert file_path == str(upload_dir / self.filename) + assert (upload_dir / self.filename).exists() + + def test_delete_file(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + self.Storage.create_container() + contents, azure_file_path = self.Storage.upload_file( + io.BytesIO(self.file_content), self.filename + ) + assert (upload_dir / self.filename).exists() + self.Storage.delete_file(azure_file_path) + assert not (upload_dir / self.filename).exists() + blob = self.Storage.blob_service_client.get_blob_client( + container=self.Storage.container_name, blob=self.filename + ) + with pytest.raises(Exception): + blob.download_blob().readall() + + def test_delete_all_files(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + self.Storage.create_container() + self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) + blob = self.Storage.blob_service_client.get_blob_client( + container=self.Storage.container_name, blob=self.filename + ) + assert self.file_content == blob.download_blob().readall() + assert (upload_dir / self.filename).exists() + self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra) + blob = self.Storage.blob_service_client.get_blob_client( + container=self.Storage.container_name, blob=self.filename_extra + ) + assert self.file_content == blob.download_blob().readall() + assert (upload_dir / self.filename_extra).exists() + + self.Storage.delete_all_files() + assert not (upload_dir / self.filename).exists() + assert not (upload_dir / self.filename_extra).exists() + blob = self.Storage.blob_service_client.get_blob_client( + container=self.Storage.container_name, blob=self.filename + ) + with pytest.raises(Exception): + blob.download_blob().readall() + blob = self.Storage.blob_service_client.get_blob_client( + container=self.Storage.container_name, blob=self.filename_extra + ) + with pytest.raises(Exception): + blob.download_blob().readall() \ No newline at end of file