mirror of
https://github.com/open-webui/open-webui
synced 2025-06-08 15:37:22 +00:00
test: add unit tests for database connection wrappers
- Implemented tests for PostgreSQL and SQLite connection handling in the wrappers module. - Covered scenarios for schema handling, existing options, and unsupported database types. - Included logging verification for successful connections.
This commit is contained in:
hobart
parent
54fe1d7f3b
commit
b203a72c61
246
backend/open_webui/test/apps/webui/internal/test_wrappers.py
Normal file
246
backend/open_webui/test/apps/webui/internal/test_wrappers.py
Normal file
@ -0,0 +1,246 @@
|
||||
import pytest
|
||||
import os
|
||||
from unittest.mock import patch, MagicMock, call
|
||||
from peewee import PostgresqlDatabase, SqliteDatabase
|
||||
|
||||
from open_webui.internal.wrappers import register_connection
|
||||
|
||||
|
||||
class TestWrappers:
|
||||
"""Test cases for the wrappers.py module, focusing on DATABASE_SCHEMA functionality."""
|
||||
|
||||
def test_register_connection_postgresql_with_schema(self):
|
||||
"""Test PostgreSQL connection when DATABASE_SCHEMA is set."""
|
||||
with patch('open_webui.internal.wrappers.DATABASE_SCHEMA', 'test_schema'):
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect, \
|
||||
patch('open_webui.internal.wrappers.parse') as mock_parse, \
|
||||
patch('open_webui.internal.wrappers.ReconnectingPostgresqlDatabase') as mock_db_class:
|
||||
|
||||
# Mock PostgreSQL database
|
||||
mock_pg_db = MagicMock(spec=PostgresqlDatabase)
|
||||
mock_connect.return_value = mock_pg_db
|
||||
|
||||
# Mock parse result without existing options
|
||||
mock_parse.return_value = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432
|
||||
}
|
||||
|
||||
# Mock database instance
|
||||
mock_db_instance = MagicMock()
|
||||
mock_db_class.return_value = mock_db_instance
|
||||
|
||||
# Test
|
||||
db_url = "postgresql://testuser:testpass@localhost:5432/testdb"
|
||||
result = register_connection(db_url)
|
||||
|
||||
# Assertions
|
||||
mock_connect.assert_called_once_with(db_url, unquote_password=True)
|
||||
mock_parse.assert_called_once_with(db_url, unquote_password=True)
|
||||
|
||||
# Check that options were added with schema
|
||||
expected_connection = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'options': '-c search_path=test_schema,public'
|
||||
}
|
||||
mock_db_class.assert_called_once_with(**expected_connection)
|
||||
mock_db_instance.connect.assert_called_once_with(reuse_if_open=True)
|
||||
|
||||
# Verify the returned database is the mock instance
|
||||
assert result == mock_db_instance
|
||||
|
||||
def test_register_connection_postgresql_with_schema_existing_options(self):
|
||||
"""Test PostgreSQL connection when DATABASE_SCHEMA is set and options already exist."""
|
||||
with patch('open_webui.internal.wrappers.DATABASE_SCHEMA', 'test_schema'):
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect, \
|
||||
patch('open_webui.internal.wrappers.parse') as mock_parse, \
|
||||
patch('open_webui.internal.wrappers.ReconnectingPostgresqlDatabase') as mock_db_class:
|
||||
|
||||
# Mock PostgreSQL database
|
||||
mock_pg_db = MagicMock(spec=PostgresqlDatabase)
|
||||
mock_connect.return_value = mock_pg_db
|
||||
|
||||
# Mock parse result with existing options
|
||||
mock_parse.return_value = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'options': '-c some_existing_option=value'
|
||||
}
|
||||
|
||||
# Mock database instance
|
||||
mock_db_instance = MagicMock()
|
||||
mock_db_class.return_value = mock_db_instance
|
||||
|
||||
# Test
|
||||
db_url = "postgresql://testuser:testpass@localhost:5432/testdb"
|
||||
result = register_connection(db_url)
|
||||
|
||||
# Check that schema was appended to existing options
|
||||
expected_connection = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'options': '-c some_existing_option=value -c search_path=test_schema,public'
|
||||
}
|
||||
mock_db_class.assert_called_once_with(**expected_connection)
|
||||
|
||||
def test_register_connection_postgresql_without_schema(self):
|
||||
"""Test PostgreSQL connection when DATABASE_SCHEMA is not set."""
|
||||
with patch('open_webui.internal.wrappers.DATABASE_SCHEMA', None):
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect, \
|
||||
patch('open_webui.internal.wrappers.parse') as mock_parse, \
|
||||
patch('open_webui.internal.wrappers.ReconnectingPostgresqlDatabase') as mock_db_class:
|
||||
|
||||
# Mock PostgreSQL database
|
||||
mock_pg_db = MagicMock(spec=PostgresqlDatabase)
|
||||
mock_connect.return_value = mock_pg_db
|
||||
|
||||
# Mock parse result
|
||||
mock_parse.return_value = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432
|
||||
}
|
||||
|
||||
# Mock database instance
|
||||
mock_db_instance = MagicMock()
|
||||
mock_db_class.return_value = mock_db_instance
|
||||
|
||||
# Test
|
||||
db_url = "postgresql://testuser:testpass@localhost:5432/testdb"
|
||||
result = register_connection(db_url)
|
||||
|
||||
# Check that no schema options were added
|
||||
expected_connection = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432
|
||||
}
|
||||
mock_db_class.assert_called_once_with(**expected_connection)
|
||||
|
||||
def test_register_connection_postgresql_empty_schema(self):
|
||||
"""Test PostgreSQL connection when DATABASE_SCHEMA is empty string."""
|
||||
with patch('open_webui.internal.wrappers.DATABASE_SCHEMA', ''):
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect, \
|
||||
patch('open_webui.internal.wrappers.parse') as mock_parse, \
|
||||
patch('open_webui.internal.wrappers.ReconnectingPostgresqlDatabase') as mock_db_class:
|
||||
|
||||
# Mock PostgreSQL database
|
||||
mock_pg_db = MagicMock(spec=PostgresqlDatabase)
|
||||
mock_connect.return_value = mock_pg_db
|
||||
|
||||
# Mock parse result
|
||||
mock_parse.return_value = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432
|
||||
}
|
||||
|
||||
# Mock database instance
|
||||
mock_db_instance = MagicMock()
|
||||
mock_db_class.return_value = mock_db_instance
|
||||
|
||||
# Test
|
||||
db_url = "postgresql://testuser:testpass@localhost:5432/testdb"
|
||||
result = register_connection(db_url)
|
||||
|
||||
# Check that no schema options were added (empty string is falsy)
|
||||
expected_connection = {
|
||||
'database': 'testdb',
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'host': 'localhost',
|
||||
'port': 5432
|
||||
}
|
||||
mock_db_class.assert_called_once_with(**expected_connection)
|
||||
|
||||
def test_register_connection_sqlite(self):
|
||||
"""Test SQLite connection (should not be affected by DATABASE_SCHEMA)."""
|
||||
with patch('open_webui.internal.wrappers.DATABASE_SCHEMA', 'test_schema'):
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect:
|
||||
|
||||
# Mock SQLite database
|
||||
mock_sqlite_db = MagicMock(spec=SqliteDatabase)
|
||||
mock_connect.return_value = mock_sqlite_db
|
||||
|
||||
# Test
|
||||
db_url = "sqlite:///test.db"
|
||||
result = register_connection(db_url)
|
||||
|
||||
# Assertions
|
||||
mock_connect.assert_called_once_with(db_url, unquote_password=True)
|
||||
|
||||
# Verify SQLite database properties are set
|
||||
assert mock_sqlite_db.autoconnect is True
|
||||
assert mock_sqlite_db.reuse_if_open is True
|
||||
|
||||
# Verify the returned database is the mock instance
|
||||
assert result == mock_sqlite_db
|
||||
|
||||
def test_register_connection_unsupported_database(self):
|
||||
"""Test that unsupported database types raise ValueError."""
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect:
|
||||
|
||||
# Mock an unsupported database type
|
||||
mock_unsupported_db = MagicMock()
|
||||
# Make it not an instance of PostgresqlDatabase or SqliteDatabase
|
||||
mock_unsupported_db.__class__ = type('UnsupportedDB', (), {})
|
||||
mock_connect.return_value = mock_unsupported_db
|
||||
|
||||
# Test
|
||||
db_url = "mysql://user:pass@localhost/db"
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported database connection"):
|
||||
register_connection(db_url)
|
||||
|
||||
@patch('open_webui.internal.wrappers.log')
|
||||
def test_register_connection_postgresql_logging(self, mock_log):
|
||||
"""Test that PostgreSQL connection logs the correct message."""
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect, \
|
||||
patch('open_webui.internal.wrappers.parse') as mock_parse, \
|
||||
patch('open_webui.internal.wrappers.ReconnectingPostgresqlDatabase') as mock_db_class:
|
||||
|
||||
# Mock PostgreSQL database
|
||||
mock_pg_db = MagicMock(spec=PostgresqlDatabase)
|
||||
mock_connect.return_value = mock_pg_db
|
||||
mock_parse.return_value = {'database': 'test'}
|
||||
mock_db_class.return_value = MagicMock()
|
||||
|
||||
# Test
|
||||
register_connection("postgresql://user:pass@localhost/db")
|
||||
|
||||
# Check logging
|
||||
mock_log.info.assert_called_with("Connected to PostgreSQL database")
|
||||
|
||||
@patch('open_webui.internal.wrappers.log')
|
||||
def test_register_connection_sqlite_logging(self, mock_log):
|
||||
"""Test that SQLite connection logs the correct message."""
|
||||
with patch('open_webui.internal.wrappers.connect') as mock_connect:
|
||||
|
||||
# Mock SQLite database
|
||||
mock_sqlite_db = MagicMock(spec=SqliteDatabase)
|
||||
mock_connect.return_value = mock_sqlite_db
|
||||
|
||||
# Test
|
||||
register_connection("sqlite:///test.db")
|
||||
|
||||
# Check logging
|
||||
mock_log.info.assert_called_with("Connected to SQLite database")
|
||||
Loading…
Reference in New Issue
Block a user