mirror of
https://github.com/open-webui/open-webui.git
synced 2026-07-12 05:25:19 +02:00
chore: remove unused legacy test suite
This commit is contained in:
Timothy Jaeryang Baek
@@ -1,196 +0,0 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestAuths(AbstractPostgresTest):
|
||||
BASE_PATH = '/api/v1/auths'
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.models.auths import Auths
|
||||
from open_webui.models.users import Users
|
||||
|
||||
cls.users = Users
|
||||
cls.auths = Auths
|
||||
|
||||
def test_get_session_user(self):
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.get(self.create_url(''))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
'id': '1',
|
||||
'name': 'John Doe',
|
||||
'email': 'john.doe@openwebui.com',
|
||||
'role': 'user',
|
||||
'profile_image_url': '/user.png',
|
||||
}
|
||||
|
||||
def test_update_profile(self):
|
||||
from open_webui.utils.auth import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password=get_password_hash('old_password'),
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='user',
|
||||
)
|
||||
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/update/profile'),
|
||||
json={'name': 'John Doe 2', 'profile_image_url': '/user2.png'},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
db_user = self.users.get_user_by_id(user.id)
|
||||
assert db_user.name == 'John Doe 2'
|
||||
assert db_user.profile_image_url == '/user2.png'
|
||||
|
||||
def test_update_password(self):
|
||||
from open_webui.utils.auth import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password=get_password_hash('old_password'),
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='user',
|
||||
)
|
||||
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/update/password'),
|
||||
json={'password': 'old_password', 'new_password': 'new_password'},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
old_auth = self.auths.authenticate_user('john.doe@openwebui.com', 'old_password')
|
||||
assert old_auth is None
|
||||
new_auth = self.auths.authenticate_user('john.doe@openwebui.com', 'new_password')
|
||||
assert new_auth is not None
|
||||
|
||||
def test_signin(self):
|
||||
from open_webui.utils.auth import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password=get_password_hash('password'),
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='user',
|
||||
)
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/signin'),
|
||||
json={'email': 'john.doe@openwebui.com', 'password': 'password'},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data['id'] == user.id
|
||||
assert data['name'] == 'John Doe'
|
||||
assert data['email'] == 'john.doe@openwebui.com'
|
||||
assert data['role'] == 'user'
|
||||
assert data['profile_image_url'] == '/user.png'
|
||||
assert data['token'] is not None and len(data['token']) > 0
|
||||
assert data['token_type'] == 'Bearer'
|
||||
|
||||
def test_signup(self):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/signup'),
|
||||
json={
|
||||
'name': 'John Doe',
|
||||
'email': 'john.doe@openwebui.com',
|
||||
'password': 'password',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data['id'] is not None and len(data['id']) > 0
|
||||
assert data['name'] == 'John Doe'
|
||||
assert data['email'] == 'john.doe@openwebui.com'
|
||||
assert data['role'] in ['admin', 'user', 'pending']
|
||||
assert data['profile_image_url'] == '/user.png'
|
||||
assert data['token'] is not None and len(data['token']) > 0
|
||||
assert data['token_type'] == 'Bearer'
|
||||
|
||||
def test_add_user(self):
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/add'),
|
||||
json={
|
||||
'name': 'John Doe 2',
|
||||
'email': 'john.doe2@openwebui.com',
|
||||
'password': 'password2',
|
||||
'role': 'admin',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data['id'] is not None and len(data['id']) > 0
|
||||
assert data['name'] == 'John Doe 2'
|
||||
assert data['email'] == 'john.doe2@openwebui.com'
|
||||
assert data['role'] == 'admin'
|
||||
assert data['profile_image_url'] == '/user.png'
|
||||
assert data['token'] is not None and len(data['token']) > 0
|
||||
assert data['token_type'] == 'Bearer'
|
||||
|
||||
def test_get_admin_details(self):
|
||||
self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password='password',
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='admin',
|
||||
)
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.get(self.create_url('/admin/details'))
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
'name': 'John Doe',
|
||||
'email': 'john.doe@openwebui.com',
|
||||
}
|
||||
|
||||
def test_create_api_key_(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password='password',
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='admin',
|
||||
)
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(self.create_url('/api_key'))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data['api_key'] is not None
|
||||
assert len(data['api_key']) > 0
|
||||
|
||||
def test_delete_api_key(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password='password',
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='admin',
|
||||
)
|
||||
self.users.update_user_api_key_by_id(user.id, 'abc')
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.delete(self.create_url('/api_key'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == True
|
||||
db_user = self.users.get_user_by_id(user.id)
|
||||
assert db_user.api_key is None
|
||||
|
||||
def test_get_api_key(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email='john.doe@openwebui.com',
|
||||
password='password',
|
||||
name='John Doe',
|
||||
profile_image_url='/user.png',
|
||||
role='admin',
|
||||
)
|
||||
self.users.update_user_api_key_by_id(user.id, 'abc')
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.get(self.create_url('/api_key'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {'api_key': 'abc'}
|
||||
@@ -1,57 +0,0 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestModels(AbstractPostgresTest):
|
||||
BASE_PATH = '/api/v1/models'
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.models.models import Model
|
||||
|
||||
cls.models = Model
|
||||
|
||||
def test_models(self):
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.get(self.create_url('/'))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/add'),
|
||||
json={
|
||||
'id': 'my-model',
|
||||
'base_model_id': 'base-model-id',
|
||||
'name': 'Hello World',
|
||||
'meta': {
|
||||
'profile_image_url': '/static/favicon.png',
|
||||
'description': 'description',
|
||||
'capabilities': None,
|
||||
'model_config': {},
|
||||
},
|
||||
'params': {},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.get(self.create_url('/'))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.get(self.create_url(query_params={'id': 'my-model'}))
|
||||
assert response.status_code == 200
|
||||
data = response.json()[0]
|
||||
assert data['id'] == 'my-model'
|
||||
assert data['name'] == 'Hello World'
|
||||
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.delete(self.create_url('/delete?id=my-model'))
|
||||
assert response.status_code == 200
|
||||
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.get(self.create_url('/'))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
@@ -1,165 +0,0 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
def _get_user_by_id(data, param):
|
||||
return next((item for item in data if item['id'] == param), None)
|
||||
|
||||
|
||||
def _assert_user(data, id, **kwargs):
|
||||
user = _get_user_by_id(data, id)
|
||||
assert user is not None
|
||||
comparison_data = {
|
||||
'name': f'user {id}',
|
||||
'email': f'user{id}@openwebui.com',
|
||||
'profile_image_url': f'/api/v1/users/{id}/profile/image',
|
||||
'role': 'user',
|
||||
**kwargs,
|
||||
}
|
||||
for key, value in comparison_data.items():
|
||||
assert user[key] == value
|
||||
|
||||
|
||||
class TestUsers(AbstractPostgresTest):
|
||||
BASE_PATH = '/api/v1/users'
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.models.users import Users
|
||||
|
||||
cls.users = Users
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
self.users.insert_new_user(
|
||||
id='1',
|
||||
name='user 1',
|
||||
email='user1@openwebui.com',
|
||||
profile_image_url='/user1.png',
|
||||
role='user',
|
||||
)
|
||||
self.users.insert_new_user(
|
||||
id='2',
|
||||
name='user 2',
|
||||
email='user2@openwebui.com',
|
||||
profile_image_url='/user2.png',
|
||||
role='user',
|
||||
)
|
||||
|
||||
def test_users(self):
|
||||
# Get all users
|
||||
with mock_webui_user(id='3'):
|
||||
response = self.fast_api_client.get(self.create_url(''))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, '1')
|
||||
_assert_user(data, '2')
|
||||
|
||||
# update role
|
||||
with mock_webui_user(id='3'):
|
||||
response = self.fast_api_client.post(self.create_url('/update/role'), json={'id': '2', 'role': 'admin'})
|
||||
assert response.status_code == 200
|
||||
_assert_user([response.json()], '2', role='admin')
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id='3'):
|
||||
response = self.fast_api_client.get(self.create_url(''))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, '1')
|
||||
_assert_user(data, '2', role='admin')
|
||||
|
||||
# Get (empty) user settings
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.get(self.create_url('/user/settings'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is None
|
||||
|
||||
# Update user settings
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/user/settings/update'),
|
||||
json={
|
||||
'ui': {'attr1': 'value1', 'attr2': 'value2'},
|
||||
'model_config': {'attr3': 'value3', 'attr4': 'value4'},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get user settings
|
||||
with mock_webui_user(id='2'):
|
||||
response = self.fast_api_client.get(self.create_url('/user/settings'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
'ui': {'attr1': 'value1', 'attr2': 'value2'},
|
||||
'model_config': {'attr3': 'value3', 'attr4': 'value4'},
|
||||
}
|
||||
|
||||
# Get (empty) user info
|
||||
with mock_webui_user(id='1'):
|
||||
response = self.fast_api_client.get(self.create_url('/user/info'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is None
|
||||
|
||||
# Update user info
|
||||
with mock_webui_user(id='1'):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/user/info/update'),
|
||||
json={'attr1': 'value1', 'attr2': 'value2'},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get user info
|
||||
with mock_webui_user(id='1'):
|
||||
response = self.fast_api_client.get(self.create_url('/user/info'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {'attr1': 'value1', 'attr2': 'value2'}
|
||||
|
||||
# Get user by id
|
||||
with mock_webui_user(id='1'):
|
||||
response = self.fast_api_client.get(self.create_url('/2'))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {'name': 'user 2', 'profile_image_url': '/user2.png'}
|
||||
|
||||
# Update user by id
|
||||
with mock_webui_user(id='1'):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url('/2/update'),
|
||||
json={
|
||||
'name': 'user 2 updated',
|
||||
'email': 'user2-updated@openwebui.com',
|
||||
'profile_image_url': '/user2-updated.png',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id='3'):
|
||||
response = self.fast_api_client.get(self.create_url(''))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, '1')
|
||||
_assert_user(
|
||||
data,
|
||||
'2',
|
||||
role='admin',
|
||||
name='user 2 updated',
|
||||
email='user2-updated@openwebui.com',
|
||||
profile_image_url=f'/api/v1/users/2/profile/image',
|
||||
)
|
||||
|
||||
# Delete user by id
|
||||
with mock_webui_user(id='1'):
|
||||
response = self.fast_api_client.delete(self.create_url('/2'))
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id='3'):
|
||||
response = self.fast_api_client.get(self.create_url(''))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
data = response.json()
|
||||
_assert_user(data, '1')
|
||||
@@ -1,406 +0,0 @@
|
||||
import io
|
||||
import os
|
||||
import boto3
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError
|
||||
from moto import mock_aws
|
||||
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, ContainerClient, BlobClient
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
def mock_upload_dir(monkeypatch, tmp_path):
|
||||
"""Fixture to monkey-patch the UPLOAD_DIR and create a temporary directory."""
|
||||
directory = tmp_path / 'uploads'
|
||||
directory.mkdir()
|
||||
monkeypatch.setattr(provider, 'UPLOAD_DIR', str(directory))
|
||||
return directory
|
||||
|
||||
|
||||
def test_imports():
|
||||
provider.StorageProvider
|
||||
provider.LocalStorageProvider
|
||||
provider.S3StorageProvider
|
||||
provider.GCSStorageProvider
|
||||
provider.AzureStorageProvider
|
||||
provider.Storage
|
||||
|
||||
|
||||
def test_get_storage_provider():
|
||||
Storage = provider.get_storage_provider('local')
|
||||
assert isinstance(Storage, provider.LocalStorageProvider)
|
||||
Storage = provider.get_storage_provider('s3')
|
||||
assert isinstance(Storage, provider.S3StorageProvider)
|
||||
Storage = provider.get_storage_provider('gcs')
|
||||
assert isinstance(Storage, provider.GCSStorageProvider)
|
||||
Storage = provider.get_storage_provider('azure')
|
||||
assert isinstance(Storage, provider.AzureStorageProvider)
|
||||
with pytest.raises(RuntimeError):
|
||||
provider.get_storage_provider('invalid')
|
||||
|
||||
|
||||
def test_class_instantiation():
|
||||
with pytest.raises(TypeError):
|
||||
provider.StorageProvider()
|
||||
with pytest.raises(TypeError):
|
||||
|
||||
class Test(provider.StorageProvider):
|
||||
pass
|
||||
|
||||
Test()
|
||||
provider.LocalStorageProvider()
|
||||
provider.S3StorageProvider()
|
||||
provider.GCSStorageProvider()
|
||||
provider.AzureStorageProvider()
|
||||
|
||||
|
||||
class TestLocalStorageProvider:
|
||||
Storage = provider.LocalStorageProvider()
|
||||
file_content = b'test content'
|
||||
file_bytesio = io.BytesIO(file_content)
|
||||
filename = 'test.txt'
|
||||
filename_extra = 'test_exyta.txt'
|
||||
file_bytesio_empty = io.BytesIO()
|
||||
|
||||
def test_upload_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
contents, file_path = self.Storage.upload_file(self.file_bytesio, self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert contents == self.file_content
|
||||
assert file_path == str(upload_dir / self.filename)
|
||||
with pytest.raises(ValueError):
|
||||
self.Storage.upload_file(self.file_bytesio_empty, self.filename)
|
||||
|
||||
def test_get_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
file_path = str(upload_dir / self.filename)
|
||||
file_path_return = self.Storage.get_file(file_path)
|
||||
assert file_path == file_path_return
|
||||
|
||||
def test_delete_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
(upload_dir / self.filename).write_bytes(self.file_content)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
file_path = str(upload_dir / self.filename)
|
||||
self.Storage.delete_file(file_path)
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
|
||||
def test_delete_all_files(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
(upload_dir / self.filename).write_bytes(self.file_content)
|
||||
(upload_dir / self.filename_extra).write_bytes(self.file_content)
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
|
||||
|
||||
@mock_aws
|
||||
class TestS3StorageProvider:
|
||||
def __init__(self):
|
||||
self.Storage = provider.S3StorageProvider()
|
||||
self.Storage.bucket_name = 'my-bucket'
|
||||
self.s3_client = boto3.resource('s3', region_name='us-east-1')
|
||||
self.file_content = b'test content'
|
||||
self.filename = 'test.txt'
|
||||
self.filename_extra = 'test_exyta.txt'
|
||||
self.file_bytesio_empty = io.BytesIO()
|
||||
super().__init__()
|
||||
|
||||
def test_upload_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# S3 checks
|
||||
with pytest.raises(Exception):
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
contents, s3_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
object = self.s3_client.Object(self.Storage.bucket_name, self.filename)
|
||||
assert self.file_content == object.get()['Body'].read()
|
||||
# local checks
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert contents == self.file_content
|
||||
assert s3_file_path == 's3://' + self.Storage.bucket_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):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
contents, s3_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
file_path = self.Storage.get_file(s3_file_path)
|
||||
assert file_path == str(upload_dir / self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
|
||||
def test_delete_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
contents, s3_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
self.Storage.delete_file(s3_file_path)
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
with pytest.raises(ClientError) as exc:
|
||||
self.s3_client.Object(self.Storage.bucket_name, self.filename).load()
|
||||
error = exc.value.response['Error']
|
||||
assert error['Code'] == '404'
|
||||
assert error['Message'] == 'Not Found'
|
||||
|
||||
def test_delete_all_files(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# create 2 files
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
object = self.s3_client.Object(self.Storage.bucket_name, self.filename)
|
||||
assert self.file_content == object.get()['Body'].read()
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra)
|
||||
object = self.s3_client.Object(self.Storage.bucket_name, self.filename_extra)
|
||||
assert self.file_content == object.get()['Body'].read()
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
with pytest.raises(ClientError) as exc:
|
||||
self.s3_client.Object(self.Storage.bucket_name, self.filename).load()
|
||||
error = exc.value.response['Error']
|
||||
assert error['Code'] == '404'
|
||||
assert error['Message'] == 'Not Found'
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
with pytest.raises(ClientError) as exc:
|
||||
self.s3_client.Object(self.Storage.bucket_name, self.filename_extra).load()
|
||||
error = exc.value.response['Error']
|
||||
assert error['Code'] == '404'
|
||||
assert error['Message'] == 'Not Found'
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
|
||||
def test_init_without_credentials(self, monkeypatch):
|
||||
"""Test that S3StorageProvider can initialize without explicit credentials."""
|
||||
# Temporarily unset the environment variables
|
||||
monkeypatch.setattr(provider, 'S3_ACCESS_KEY_ID', None)
|
||||
monkeypatch.setattr(provider, 'S3_SECRET_ACCESS_KEY', None)
|
||||
|
||||
# Should not raise an exception
|
||||
storage = provider.S3StorageProvider()
|
||||
assert storage.s3_client is not None
|
||||
assert storage.bucket_name == provider.S3_BUCKET_NAME
|
||||
|
||||
|
||||
class TestGCSStorageProvider:
|
||||
Storage = provider.GCSStorageProvider()
|
||||
Storage.bucket_name = 'my-bucket'
|
||||
file_content = b'test content'
|
||||
filename = 'test.txt'
|
||||
filename_extra = 'test_exyta.txt'
|
||||
file_bytesio_empty = io.BytesIO()
|
||||
|
||||
@pytest.fixture(scope='class')
|
||||
def setup(self):
|
||||
host, port = 'localhost', 9023
|
||||
|
||||
server = create_server(host, port, in_memory=True)
|
||||
server.start()
|
||||
os.environ['STORAGE_EMULATOR_HOST'] = f'http://{host}:{port}'
|
||||
|
||||
gcs_client = storage.Client()
|
||||
bucket = gcs_client.bucket(self.Storage.bucket_name)
|
||||
bucket.create()
|
||||
self.Storage.gcs_client, self.Storage.bucket = gcs_client, bucket
|
||||
yield
|
||||
bucket.delete(force=True)
|
||||
server.stop()
|
||||
|
||||
def test_upload_file(self, monkeypatch, tmp_path, setup):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# catch error if bucket does not exist
|
||||
with pytest.raises(Exception):
|
||||
self.Storage.bucket = monkeypatch(self.Storage, 'bucket', None)
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
contents, gcs_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
object = self.Storage.bucket.get_blob(self.filename)
|
||||
assert self.file_content == object.download_as_bytes()
|
||||
# local checks
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert contents == self.file_content
|
||||
assert gcs_file_path == 'gs://' + self.Storage.bucket_name + '/' + self.filename
|
||||
# test error if file is empty
|
||||
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)
|
||||
contents, gcs_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
file_path = self.Storage.get_file(gcs_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)
|
||||
contents, gcs_file_path = self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
# ensure that local directory has the uploaded file as well
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert self.Storage.bucket.get_blob(self.filename).name == self.filename
|
||||
self.Storage.delete_file(gcs_file_path)
|
||||
# check that deleting file from gcs will delete the local file as well
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert self.Storage.bucket.get_blob(self.filename) == None
|
||||
|
||||
def test_delete_all_files(self, monkeypatch, tmp_path, setup):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# create 2 files
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
object = self.Storage.bucket.get_blob(self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert self.Storage.bucket.get_blob(self.filename).name == self.filename
|
||||
assert self.file_content == object.download_as_bytes()
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra)
|
||||
object = self.Storage.bucket.get_blob(self.filename_extra)
|
||||
assert (upload_dir / self.filename_extra).exists()
|
||||
assert (upload_dir / self.filename_extra).read_bytes() == self.file_content
|
||||
assert self.Storage.bucket.get_blob(self.filename_extra).name == self.filename_extra
|
||||
assert self.file_content == object.download_as_bytes()
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
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):
|
||||
super().__init__()
|
||||
|
||||
@pytest.fixture(scope='class')
|
||||
def setup_storage(self, monkeypatch):
|
||||
# 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,
|
||||
'ContainerClient',
|
||||
lambda *args, **kwargs: mock_container_client,
|
||||
)
|
||||
monkeypatch.setattr(azure.storage.blob, 'BlobClient', lambda *args, **kwargs: mock_blob_client)
|
||||
|
||||
self.Storage = provider.AzureStorageProvider()
|
||||
self.Storage.endpoint = 'https://myaccount.blob.core.windows.net'
|
||||
self.Storage.container_name = 'my-container'
|
||||
self.file_content = b'test content'
|
||||
self.filename = 'test.txt'
|
||||
self.filename_extra = 'test_extra.txt'
|
||||
self.file_bytesio_empty = io.BytesIO()
|
||||
|
||||
# Apply mocks to the Storage instance
|
||||
self.Storage.blob_service_client = mock_blob_service_client
|
||||
self.Storage.container_client = mock_container_client
|
||||
|
||||
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)
|
||||
|
||||
# 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
|
||||
)
|
||||
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()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
self.Storage.upload_file(self.file_bytesio_empty, self.filename)
|
||||
|
||||
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_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):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
self.Storage.create_container()
|
||||
|
||||
# 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
|
||||
|
||||
file_url = f'https://myaccount.blob.core.windows.net/{self.Storage.container_name}/{self.filename}'
|
||||
self.Storage.delete_file(file_url)
|
||||
|
||||
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):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
self.Storage.create_container()
|
||||
|
||||
# Mock file uploads
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra)
|
||||
|
||||
# Mock listing and deletion behavior
|
||||
self.Storage.container_client.list_blobs.return_value = [
|
||||
{'name': self.filename},
|
||||
{'name': self.filename_extra},
|
||||
]
|
||||
self.Storage.container_client.get_blob_client().delete_blob.return_value = None
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
|
||||
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):
|
||||
self.Storage.create_container()
|
||||
|
||||
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(file_url)
|
||||
@@ -1,781 +0,0 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch, AsyncMock
|
||||
import redis
|
||||
from open_webui.utils.redis import (
|
||||
SentinelRedisProxy,
|
||||
parse_redis_service_url,
|
||||
get_redis_connection,
|
||||
get_sentinels_from_env,
|
||||
MAX_RETRY_COUNT,
|
||||
)
|
||||
import inspect
|
||||
|
||||
|
||||
class TestSentinelRedisProxy:
|
||||
"""Test Redis Sentinel failover functionality"""
|
||||
|
||||
def test_parse_redis_service_url_valid(self):
|
||||
"""Test parsing valid Redis service URL"""
|
||||
url = 'redis://user:pass@mymaster:6379/0'
|
||||
result = parse_redis_service_url(url)
|
||||
|
||||
assert result['username'] == 'user'
|
||||
assert result['password'] == 'pass'
|
||||
assert result['service'] == 'mymaster'
|
||||
assert result['port'] == 6379
|
||||
assert result['db'] == 0
|
||||
|
||||
def test_parse_redis_service_url_defaults(self):
|
||||
"""Test parsing Redis service URL with defaults"""
|
||||
url = 'redis://mymaster'
|
||||
result = parse_redis_service_url(url)
|
||||
|
||||
assert result['username'] is None
|
||||
assert result['password'] is None
|
||||
assert result['service'] == 'mymaster'
|
||||
assert result['port'] == 6379
|
||||
assert result['db'] == 0
|
||||
|
||||
def test_parse_redis_service_url_invalid_scheme(self):
|
||||
"""Test parsing invalid URL scheme"""
|
||||
with pytest.raises(ValueError, match='Invalid Redis URL scheme'):
|
||||
parse_redis_service_url('http://invalid')
|
||||
|
||||
def test_get_sentinels_from_env(self):
|
||||
"""Test parsing sentinel hosts from environment"""
|
||||
hosts = 'sentinel1,sentinel2,sentinel3'
|
||||
port = '26379'
|
||||
|
||||
result = get_sentinels_from_env(hosts, port)
|
||||
expected = [('sentinel1', 26379), ('sentinel2', 26379), ('sentinel3', 26379)]
|
||||
|
||||
assert result == expected
|
||||
|
||||
def test_get_sentinels_from_env_empty(self):
|
||||
"""Test empty sentinel hosts"""
|
||||
result = get_sentinels_from_env(None, '26379')
|
||||
assert result == []
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_sentinel_redis_proxy_sync_success(self, mock_sentinel_class):
|
||||
"""Test successful sync operation with SentinelRedisProxy"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_master.get.return_value = 'test_value'
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test attribute access
|
||||
get_method = proxy.__getattr__('get')
|
||||
result = get_method('test_key')
|
||||
|
||||
assert result == 'test_value'
|
||||
mock_sentinel.master_for.assert_called_with('mymaster')
|
||||
mock_master.get.assert_called_with('test_key')
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_sentinel_redis_proxy_async_success(self, mock_sentinel_class):
|
||||
"""Test successful async operation with SentinelRedisProxy"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_master.get = AsyncMock(return_value='test_value')
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test async attribute access
|
||||
get_method = proxy.__getattr__('get')
|
||||
result = await get_method('test_key')
|
||||
|
||||
assert result == 'test_value'
|
||||
mock_sentinel.master_for.assert_called_with('mymaster')
|
||||
mock_master.get.assert_called_with('test_key')
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_sentinel_redis_proxy_failover_retry(self, mock_sentinel_class):
|
||||
"""Test retry mechanism during failover"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# First call fails, second succeeds
|
||||
mock_master.get.side_effect = [
|
||||
redis.exceptions.ConnectionError('Master down'),
|
||||
'test_value',
|
||||
]
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
get_method = proxy.__getattr__('get')
|
||||
result = get_method('test_key')
|
||||
|
||||
assert result == 'test_value'
|
||||
assert mock_master.get.call_count == 2
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_sentinel_redis_proxy_max_retries_exceeded(self, mock_sentinel_class):
|
||||
"""Test failure after max retries exceeded"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# All calls fail
|
||||
mock_master.get.side_effect = redis.exceptions.ConnectionError('Master down')
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
get_method = proxy.__getattr__('get')
|
||||
|
||||
with pytest.raises(redis.exceptions.ConnectionError):
|
||||
get_method('test_key')
|
||||
|
||||
assert mock_master.get.call_count == MAX_RETRY_COUNT
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_sentinel_redis_proxy_readonly_error_retry(self, mock_sentinel_class):
|
||||
"""Test retry on ReadOnlyError"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# First call gets ReadOnlyError (old master), second succeeds (new master)
|
||||
mock_master.get.side_effect = [
|
||||
redis.exceptions.ReadOnlyError('Read only'),
|
||||
'test_value',
|
||||
]
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
get_method = proxy.__getattr__('get')
|
||||
result = get_method('test_key')
|
||||
|
||||
assert result == 'test_value'
|
||||
assert mock_master.get.call_count == 2
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_sentinel_redis_proxy_factory_methods(self, mock_sentinel_class):
|
||||
"""Test factory methods are passed through directly"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_pipeline = Mock()
|
||||
mock_master.pipeline.return_value = mock_pipeline
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Factory methods should be passed through without wrapping
|
||||
pipeline_method = proxy.__getattr__('pipeline')
|
||||
result = pipeline_method()
|
||||
|
||||
assert result == mock_pipeline
|
||||
mock_master.pipeline.assert_called_once()
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@patch('redis.from_url')
|
||||
def test_get_redis_connection_with_sentinel(self, mock_from_url, mock_sentinel_class):
|
||||
"""Test getting Redis connection with Sentinel"""
|
||||
mock_sentinel = Mock()
|
||||
mock_sentinel_class.return_value = mock_sentinel
|
||||
|
||||
sentinels = [('sentinel1', 26379), ('sentinel2', 26379)]
|
||||
redis_url = 'redis://user:pass@mymaster:6379/0'
|
||||
|
||||
result = get_redis_connection(redis_url=redis_url, redis_sentinels=sentinels, async_mode=False)
|
||||
|
||||
assert isinstance(result, SentinelRedisProxy)
|
||||
mock_sentinel_class.assert_called_once()
|
||||
mock_from_url.assert_not_called()
|
||||
|
||||
@patch('redis.Redis.from_url')
|
||||
def test_get_redis_connection_without_sentinel(self, mock_from_url):
|
||||
"""Test getting Redis connection without Sentinel"""
|
||||
mock_redis = Mock()
|
||||
mock_from_url.return_value = mock_redis
|
||||
|
||||
redis_url = 'redis://localhost:6379/0'
|
||||
|
||||
result = get_redis_connection(redis_url=redis_url, redis_sentinels=None, async_mode=False)
|
||||
|
||||
assert result == mock_redis
|
||||
mock_from_url.assert_called_once_with(redis_url, decode_responses=True)
|
||||
|
||||
@patch('redis.asyncio.from_url')
|
||||
def test_get_redis_connection_without_sentinel_async(self, mock_from_url):
|
||||
"""Test getting async Redis connection without Sentinel"""
|
||||
mock_redis = Mock()
|
||||
mock_from_url.return_value = mock_redis
|
||||
|
||||
redis_url = 'redis://localhost:6379/0'
|
||||
|
||||
result = get_redis_connection(redis_url=redis_url, redis_sentinels=None, async_mode=True)
|
||||
|
||||
assert result == mock_redis
|
||||
mock_from_url.assert_called_once_with(redis_url, decode_responses=True)
|
||||
|
||||
|
||||
class TestSentinelRedisProxyCommands:
|
||||
"""Test Redis commands through SentinelRedisProxy"""
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_hash_commands_sync(self, mock_sentinel_class):
|
||||
"""Test Redis hash commands in sync mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock hash command responses
|
||||
mock_master.hset.return_value = 1
|
||||
mock_master.hget.return_value = 'test_value'
|
||||
mock_master.hgetall.return_value = {'key1': 'value1', 'key2': 'value2'}
|
||||
mock_master.hdel.return_value = 1
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test hset
|
||||
hset_method = proxy.__getattr__('hset')
|
||||
result = hset_method('test_hash', 'field1', 'value1')
|
||||
assert result == 1
|
||||
mock_master.hset.assert_called_with('test_hash', 'field1', 'value1')
|
||||
|
||||
# Test hget
|
||||
hget_method = proxy.__getattr__('hget')
|
||||
result = hget_method('test_hash', 'field1')
|
||||
assert result == 'test_value'
|
||||
mock_master.hget.assert_called_with('test_hash', 'field1')
|
||||
|
||||
# Test hgetall
|
||||
hgetall_method = proxy.__getattr__('hgetall')
|
||||
result = hgetall_method('test_hash')
|
||||
assert result == {'key1': 'value1', 'key2': 'value2'}
|
||||
mock_master.hgetall.assert_called_with('test_hash')
|
||||
|
||||
# Test hdel
|
||||
hdel_method = proxy.__getattr__('hdel')
|
||||
result = hdel_method('test_hash', 'field1')
|
||||
assert result == 1
|
||||
mock_master.hdel.assert_called_with('test_hash', 'field1')
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_hash_commands_async(self, mock_sentinel_class):
|
||||
"""Test Redis hash commands in async mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock async hash command responses
|
||||
mock_master.hset = AsyncMock(return_value=1)
|
||||
mock_master.hget = AsyncMock(return_value='test_value')
|
||||
mock_master.hgetall = AsyncMock(return_value={'key1': 'value1', 'key2': 'value2'})
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test hset
|
||||
hset_method = proxy.__getattr__('hset')
|
||||
result = await hset_method('test_hash', 'field1', 'value1')
|
||||
assert result == 1
|
||||
mock_master.hset.assert_called_with('test_hash', 'field1', 'value1')
|
||||
|
||||
# Test hget
|
||||
hget_method = proxy.__getattr__('hget')
|
||||
result = await hget_method('test_hash', 'field1')
|
||||
assert result == 'test_value'
|
||||
mock_master.hget.assert_called_with('test_hash', 'field1')
|
||||
|
||||
# Test hgetall
|
||||
hgetall_method = proxy.__getattr__('hgetall')
|
||||
result = await hgetall_method('test_hash')
|
||||
assert result == {'key1': 'value1', 'key2': 'value2'}
|
||||
mock_master.hgetall.assert_called_with('test_hash')
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_string_commands_sync(self, mock_sentinel_class):
|
||||
"""Test Redis string commands in sync mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock string command responses
|
||||
mock_master.set.return_value = True
|
||||
mock_master.get.return_value = 'test_value'
|
||||
mock_master.delete.return_value = 1
|
||||
mock_master.exists.return_value = True
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test set
|
||||
set_method = proxy.__getattr__('set')
|
||||
result = set_method('test_key', 'test_value')
|
||||
assert result is True
|
||||
mock_master.set.assert_called_with('test_key', 'test_value')
|
||||
|
||||
# Test get
|
||||
get_method = proxy.__getattr__('get')
|
||||
result = get_method('test_key')
|
||||
assert result == 'test_value'
|
||||
mock_master.get.assert_called_with('test_key')
|
||||
|
||||
# Test delete
|
||||
delete_method = proxy.__getattr__('delete')
|
||||
result = delete_method('test_key')
|
||||
assert result == 1
|
||||
mock_master.delete.assert_called_with('test_key')
|
||||
|
||||
# Test exists
|
||||
exists_method = proxy.__getattr__('exists')
|
||||
result = exists_method('test_key')
|
||||
assert result is True
|
||||
mock_master.exists.assert_called_with('test_key')
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_list_commands_sync(self, mock_sentinel_class):
|
||||
"""Test Redis list commands in sync mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock list command responses
|
||||
mock_master.lpush.return_value = 1
|
||||
mock_master.rpop.return_value = 'test_value'
|
||||
mock_master.llen.return_value = 5
|
||||
mock_master.lrange.return_value = ['item1', 'item2', 'item3']
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test lpush
|
||||
lpush_method = proxy.__getattr__('lpush')
|
||||
result = lpush_method('test_list', 'item1')
|
||||
assert result == 1
|
||||
mock_master.lpush.assert_called_with('test_list', 'item1')
|
||||
|
||||
# Test rpop
|
||||
rpop_method = proxy.__getattr__('rpop')
|
||||
result = rpop_method('test_list')
|
||||
assert result == 'test_value'
|
||||
mock_master.rpop.assert_called_with('test_list')
|
||||
|
||||
# Test llen
|
||||
llen_method = proxy.__getattr__('llen')
|
||||
result = llen_method('test_list')
|
||||
assert result == 5
|
||||
mock_master.llen.assert_called_with('test_list')
|
||||
|
||||
# Test lrange
|
||||
lrange_method = proxy.__getattr__('lrange')
|
||||
result = lrange_method('test_list', 0, -1)
|
||||
assert result == ['item1', 'item2', 'item3']
|
||||
mock_master.lrange.assert_called_with('test_list', 0, -1)
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_pubsub_commands_sync(self, mock_sentinel_class):
|
||||
"""Test Redis pubsub commands in sync mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_pubsub = Mock()
|
||||
|
||||
# Mock pubsub responses
|
||||
mock_master.pubsub.return_value = mock_pubsub
|
||||
mock_master.publish.return_value = 1
|
||||
mock_pubsub.subscribe.return_value = None
|
||||
mock_pubsub.get_message.return_value = {'type': 'message', 'data': 'test_data'}
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test pubsub (factory method - should pass through)
|
||||
pubsub_method = proxy.__getattr__('pubsub')
|
||||
result = pubsub_method()
|
||||
assert result == mock_pubsub
|
||||
mock_master.pubsub.assert_called_once()
|
||||
|
||||
# Test publish
|
||||
publish_method = proxy.__getattr__('publish')
|
||||
result = publish_method('test_channel', 'test_message')
|
||||
assert result == 1
|
||||
mock_master.publish.assert_called_with('test_channel', 'test_message')
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_pipeline_commands_sync(self, mock_sentinel_class):
|
||||
"""Test Redis pipeline commands in sync mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_pipeline = Mock()
|
||||
|
||||
# Mock pipeline responses
|
||||
mock_master.pipeline.return_value = mock_pipeline
|
||||
mock_pipeline.set.return_value = mock_pipeline
|
||||
mock_pipeline.get.return_value = mock_pipeline
|
||||
mock_pipeline.execute.return_value = [True, 'test_value']
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test pipeline (factory method - should pass through)
|
||||
pipeline_method = proxy.__getattr__('pipeline')
|
||||
result = pipeline_method()
|
||||
assert result == mock_pipeline
|
||||
mock_master.pipeline.assert_called_once()
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_commands_with_failover_retry(self, mock_sentinel_class):
|
||||
"""Test Redis commands with failover retry mechanism"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# First call fails with connection error, second succeeds
|
||||
mock_master.hget.side_effect = [
|
||||
redis.exceptions.ConnectionError('Connection failed'),
|
||||
'recovered_value',
|
||||
]
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test hget with retry
|
||||
hget_method = proxy.__getattr__('hget')
|
||||
result = hget_method('test_hash', 'field1')
|
||||
|
||||
assert result == 'recovered_value'
|
||||
assert mock_master.hget.call_count == 2
|
||||
|
||||
# Verify both calls were made with same parameters
|
||||
expected_calls = [(('test_hash', 'field1'),), (('test_hash', 'field1'),)]
|
||||
actual_calls = [call.args for call in mock_master.hget.call_args_list]
|
||||
assert actual_calls == expected_calls
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
def test_commands_with_readonly_error_retry(self, mock_sentinel_class):
|
||||
"""Test Redis commands with ReadOnlyError retry mechanism"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# First call fails with ReadOnlyError, second succeeds
|
||||
mock_master.hset.side_effect = [
|
||||
redis.exceptions.ReadOnlyError("READONLY You can't write against a read only replica"),
|
||||
1,
|
||||
]
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=False)
|
||||
|
||||
# Test hset with retry
|
||||
hset_method = proxy.__getattr__('hset')
|
||||
result = hset_method('test_hash', 'field1', 'value1')
|
||||
|
||||
assert result == 1
|
||||
assert mock_master.hset.call_count == 2
|
||||
|
||||
# Verify both calls were made with same parameters
|
||||
expected_calls = [
|
||||
(('test_hash', 'field1', 'value1'),),
|
||||
(('test_hash', 'field1', 'value1'),),
|
||||
]
|
||||
actual_calls = [call.args for call in mock_master.hset.call_args_list]
|
||||
assert actual_calls == expected_calls
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_commands_with_failover_retry(self, mock_sentinel_class):
|
||||
"""Test async Redis commands with failover retry mechanism"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# First call fails with connection error, second succeeds
|
||||
mock_master.hget = AsyncMock(
|
||||
side_effect=[
|
||||
redis.exceptions.ConnectionError('Connection failed'),
|
||||
'recovered_value',
|
||||
]
|
||||
)
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test async hget with retry
|
||||
hget_method = proxy.__getattr__('hget')
|
||||
result = await hget_method('test_hash', 'field1')
|
||||
|
||||
assert result == 'recovered_value'
|
||||
assert mock_master.hget.call_count == 2
|
||||
|
||||
# Verify both calls were made with same parameters
|
||||
expected_calls = [(('test_hash', 'field1'),), (('test_hash', 'field1'),)]
|
||||
actual_calls = [call.args for call in mock_master.hget.call_args_list]
|
||||
assert actual_calls == expected_calls
|
||||
|
||||
|
||||
class TestSentinelRedisProxyFactoryMethods:
|
||||
"""Test Redis factory methods in async mode - these are special cases that remain sync"""
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_pubsub_factory_method_async(self, mock_sentinel_class):
|
||||
"""Test pubsub factory method in async mode - should pass through without wrapping"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_pubsub = Mock()
|
||||
|
||||
# Mock pubsub factory method
|
||||
mock_master.pubsub.return_value = mock_pubsub
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test pubsub factory method - should NOT be wrapped as async
|
||||
pubsub_method = proxy.__getattr__('pubsub')
|
||||
result = pubsub_method()
|
||||
|
||||
assert result == mock_pubsub
|
||||
mock_master.pubsub.assert_called_once()
|
||||
|
||||
# Verify it's not wrapped as async (no await needed)
|
||||
assert not inspect.iscoroutine(result)
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_factory_method_async(self, mock_sentinel_class):
|
||||
"""Test pipeline factory method in async mode - should pass through without wrapping"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_pipeline = Mock()
|
||||
|
||||
# Mock pipeline factory method
|
||||
mock_master.pipeline.return_value = mock_pipeline
|
||||
mock_pipeline.set.return_value = mock_pipeline
|
||||
mock_pipeline.get.return_value = mock_pipeline
|
||||
mock_pipeline.execute.return_value = [True, 'test_value']
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test pipeline factory method - should NOT be wrapped as async
|
||||
pipeline_method = proxy.__getattr__('pipeline')
|
||||
result = pipeline_method()
|
||||
|
||||
assert result == mock_pipeline
|
||||
mock_master.pipeline.assert_called_once()
|
||||
|
||||
# Verify it's not wrapped as async (no await needed)
|
||||
assert not inspect.iscoroutine(result)
|
||||
|
||||
# Test pipeline usage (these should also be sync)
|
||||
pipeline_result = result.set('key', 'value').get('key').execute()
|
||||
assert pipeline_result == [True, 'test_value']
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_methods_vs_regular_commands_async(self, mock_sentinel_class):
|
||||
"""Test that factory methods behave differently from regular commands in async mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock both factory method and regular command
|
||||
mock_pubsub = Mock()
|
||||
mock_master.pubsub.return_value = mock_pubsub
|
||||
mock_master.get = AsyncMock(return_value='test_value')
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test factory method - should NOT be wrapped
|
||||
pubsub_method = proxy.__getattr__('pubsub')
|
||||
pubsub_result = pubsub_method()
|
||||
|
||||
# Test regular command - should be wrapped as async
|
||||
get_method = proxy.__getattr__('get')
|
||||
get_result = get_method('test_key')
|
||||
|
||||
# Factory method returns directly
|
||||
assert pubsub_result == mock_pubsub
|
||||
assert not inspect.iscoroutine(pubsub_result)
|
||||
|
||||
# Regular command returns coroutine
|
||||
assert inspect.iscoroutine(get_result)
|
||||
|
||||
# Regular command needs await
|
||||
actual_value = await get_result
|
||||
assert actual_value == 'test_value'
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_methods_with_failover_async(self, mock_sentinel_class):
|
||||
"""Test factory methods with failover in async mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# First call fails, second succeeds
|
||||
mock_pubsub = Mock()
|
||||
mock_master.pubsub.side_effect = [
|
||||
redis.exceptions.ConnectionError('Connection failed'),
|
||||
mock_pubsub,
|
||||
]
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test pubsub factory method with failover
|
||||
pubsub_method = proxy.__getattr__('pubsub')
|
||||
result = pubsub_method()
|
||||
|
||||
assert result == mock_pubsub
|
||||
assert mock_master.pubsub.call_count == 2 # Retry happened
|
||||
|
||||
# Verify it's still not wrapped as async after retry
|
||||
assert not inspect.iscoroutine(result)
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitor_factory_method_async(self, mock_sentinel_class):
|
||||
"""Test monitor factory method in async mode - should pass through without wrapping"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_monitor = Mock()
|
||||
|
||||
# Mock monitor factory method
|
||||
mock_master.monitor.return_value = mock_monitor
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test monitor factory method - should NOT be wrapped as async
|
||||
monitor_method = proxy.__getattr__('monitor')
|
||||
result = monitor_method()
|
||||
|
||||
assert result == mock_monitor
|
||||
mock_master.monitor.assert_called_once()
|
||||
|
||||
# Verify it's not wrapped as async (no await needed)
|
||||
assert not inspect.iscoroutine(result)
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_factory_method_async(self, mock_sentinel_class):
|
||||
"""Test client factory method in async mode - should pass through without wrapping"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_client = Mock()
|
||||
|
||||
# Mock client factory method
|
||||
mock_master.client.return_value = mock_client
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test client factory method - should NOT be wrapped as async
|
||||
client_method = proxy.__getattr__('client')
|
||||
result = client_method()
|
||||
|
||||
assert result == mock_client
|
||||
mock_master.client.assert_called_once()
|
||||
|
||||
# Verify it's not wrapped as async (no await needed)
|
||||
assert not inspect.iscoroutine(result)
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_transaction_factory_method_async(self, mock_sentinel_class):
|
||||
"""Test transaction factory method in async mode - should pass through without wrapping"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
mock_transaction = Mock()
|
||||
|
||||
# Mock transaction factory method
|
||||
mock_master.transaction.return_value = mock_transaction
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test transaction factory method - should NOT be wrapped as async
|
||||
transaction_method = proxy.__getattr__('transaction')
|
||||
result = transaction_method()
|
||||
|
||||
assert result == mock_transaction
|
||||
mock_master.transaction.assert_called_once()
|
||||
|
||||
# Verify it's not wrapped as async (no await needed)
|
||||
assert not inspect.iscoroutine(result)
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_factory_methods_async(self, mock_sentinel_class):
|
||||
"""Test all factory methods in async mode - comprehensive test"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock all factory methods
|
||||
mock_objects = {
|
||||
'pipeline': Mock(),
|
||||
'pubsub': Mock(),
|
||||
'monitor': Mock(),
|
||||
'client': Mock(),
|
||||
'transaction': Mock(),
|
||||
}
|
||||
|
||||
for method_name, mock_obj in mock_objects.items():
|
||||
setattr(mock_master, method_name, Mock(return_value=mock_obj))
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Test all factory methods
|
||||
for method_name, expected_obj in mock_objects.items():
|
||||
method = proxy.__getattr__(method_name)
|
||||
result = method()
|
||||
|
||||
assert result == expected_obj
|
||||
assert not inspect.iscoroutine(result)
|
||||
getattr(mock_master, method_name).assert_called_once()
|
||||
|
||||
# Reset mock for next iteration
|
||||
getattr(mock_master, method_name).reset_mock()
|
||||
|
||||
@patch('redis.sentinel.Sentinel')
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_factory_and_regular_commands_async(self, mock_sentinel_class):
|
||||
"""Test using both factory methods and regular commands in async mode"""
|
||||
mock_sentinel = Mock()
|
||||
mock_master = Mock()
|
||||
|
||||
# Mock pipeline factory and regular commands
|
||||
mock_pipeline = Mock()
|
||||
mock_master.pipeline.return_value = mock_pipeline
|
||||
mock_pipeline.set.return_value = mock_pipeline
|
||||
mock_pipeline.get.return_value = mock_pipeline
|
||||
mock_pipeline.execute.return_value = [True, 'pipeline_value']
|
||||
|
||||
mock_master.get = AsyncMock(return_value='regular_value')
|
||||
|
||||
mock_sentinel.master_for.return_value = mock_master
|
||||
|
||||
proxy = SentinelRedisProxy(mock_sentinel, 'mymaster', async_mode=True)
|
||||
|
||||
# Use factory method (sync)
|
||||
pipeline = proxy.__getattr__('pipeline')()
|
||||
pipeline_result = pipeline.set('key1', 'value1').get('key1').execute()
|
||||
|
||||
# Use regular command (async)
|
||||
get_method = proxy.__getattr__('get')
|
||||
regular_result = await get_method('key2')
|
||||
|
||||
# Verify both work correctly
|
||||
assert pipeline_result == [True, 'pipeline_value']
|
||||
assert regular_result == 'regular_value'
|
||||
|
||||
# Verify calls
|
||||
mock_master.pipeline.assert_called_once()
|
||||
mock_master.get.assert_called_with('key2')
|
||||
Reference in New Issue
Block a user