From b203a72c61543b05b5b89c4617b893f39deb707c Mon Sep 17 00:00:00 2001 From: hobart Date: Fri, 23 May 2025 15:15:06 +0900 Subject: [PATCH] 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. --- .../test/apps/webui/internal/test_wrappers.py | 246 ++++++++++++++++++ 1 file changed, 246 insertions(+) create mode 100644 backend/open_webui/test/apps/webui/internal/test_wrappers.py diff --git a/backend/open_webui/test/apps/webui/internal/test_wrappers.py b/backend/open_webui/test/apps/webui/internal/test_wrappers.py new file mode 100644 index 000000000..3fbb50e57 --- /dev/null +++ b/backend/open_webui/test/apps/webui/internal/test_wrappers.py @@ -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")