add tests

This commit is contained in:
Chris Pietschmann 2025-02-18 14:09:00 -05:00
parent 13227fc017
commit 2f4f4c2376
2 changed files with 100 additions and 1 deletions

View File

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

View File

@ -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()