Browse Source

add tests

Chris Pietschmann 2 months ago
parent
commit
2f4f4c2376

+ 1 - 0
.github/workflows/integration-test.disabled

@@ -5,6 +5,7 @@ on:
     branches:
       - main
       - dev
+      - azure-storage
   pull_request:
     branches:
       - main

+ 99 - 1
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()