-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix(core): Fix not create the database bug
- Loading branch information
Showing
3 changed files
with
118 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,71 @@ | ||
import pytest | ||
from unittest.mock import patch, MagicMock | ||
from sqlalchemy.exc import OperationalError, SQLAlchemyError | ||
|
||
from dbgpt.app.base import _create_mysql_database | ||
|
||
|
||
@patch("sqlalchemy.create_engine") | ||
@patch("dbgpt.app.base.logger") | ||
def test_database_already_exists(mock_logger, mock_create_engine): | ||
mock_connection = MagicMock() | ||
mock_create_engine.return_value.connect.return_value.__enter__.return_value = ( | ||
mock_connection | ||
) | ||
|
||
_create_mysql_database( | ||
"test_db", "mysql+pymysql://user:password@host/test_db", True | ||
) | ||
mock_logger.info.assert_called_with("Database test_db already exists") | ||
mock_connection.execute.assert_not_called() | ||
|
||
|
||
@patch("sqlalchemy.create_engine") | ||
@patch("dbgpt.app.base.logger") | ||
def test_database_creation_success(mock_logger, mock_create_engine): | ||
# Mock the first connection failure, and the second connection success | ||
mock_create_engine.side_effect = [ | ||
MagicMock( | ||
connect=MagicMock( | ||
side_effect=OperationalError("Unknown database", None, None) | ||
) | ||
), | ||
MagicMock(), | ||
] | ||
|
||
_create_mysql_database( | ||
"test_db", "mysql+pymysql://user:password@host/test_db", True | ||
) | ||
mock_logger.info.assert_called_with("Database test_db successfully created") | ||
|
||
|
||
@patch("sqlalchemy.create_engine") | ||
@patch("dbgpt.app.base.logger") | ||
def test_database_creation_failure(mock_logger, mock_create_engine): | ||
# Mock the first connection failure, and the second connection failure with SQLAlchemyError | ||
mock_create_engine.side_effect = [ | ||
MagicMock( | ||
connect=MagicMock( | ||
side_effect=OperationalError("Unknown database", None, None) | ||
) | ||
), | ||
MagicMock(connect=MagicMock(side_effect=SQLAlchemyError("Creation failed"))), | ||
] | ||
|
||
with pytest.raises(SQLAlchemyError): | ||
_create_mysql_database( | ||
"test_db", "mysql+pymysql://user:password@host/test_db", True | ||
) | ||
mock_logger.error.assert_called_with( | ||
"Failed to create database test_db: Creation failed" | ||
) | ||
|
||
|
||
@patch("sqlalchemy.create_engine") | ||
@patch("dbgpt.app.base.logger") | ||
def test_skip_database_creation(mock_logger, mock_create_engine): | ||
_create_mysql_database( | ||
"test_db", "mysql+pymysql://user:password@host/test_db", False | ||
) | ||
mock_logger.info.assert_called_with("Skipping creation of database test_db") | ||
mock_create_engine.assert_not_called() |