Nettacker/tests/database/test_mysql.py

from unittest.mock import patch, MagicMock

from sqlalchemy.exc import SQLAlchemyError

from nettacker.config import Config
from nettacker.database.models import Base
from nettacker.database.mysql import mysql_create_database, mysql_create_tables
from tests.common import TestCase


class TestMySQLFunctions(TestCase):
    """Test cases for mysql.py functions"""

    @patch("nettacker.database.mysql.create_engine")
    def test_mysql_create_database_success(self, mock_create_engine):
        """Test successful database creation"""
        # Set up mock config
        Config.db = MagicMock()
        Config.db.as_dict.return_value = {
            "username": "test_user",
            "password": "test_pass",
            "host": "localhost",
            "port": "3306",
            "name": "test_db",
        }
        Config.db.name = "test_db"

        # Set up mock connection and execution
        mock_conn = MagicMock()
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine
        mock_engine.connect.return_value.__enter__.return_value = mock_conn

        # Mock database query results - database doesn't exist yet
        mock_conn.execute.return_value = [("mysql",), ("information_schema",)]

        # Call the function
        mysql_create_database()

        # Assertions
        mock_create_engine.assert_called_once_with(
            "mysql+pymysql://test_user:test_pass@localhost:3306"
        )

        # Check that execute was called with any text object that has the expected SQL
        call_args_list = mock_conn.execute.call_args_list
        self.assertEqual(len(call_args_list), 2)  # Two calls to execute

        # Check that the first call is SHOW DATABASES
        first_call_arg = call_args_list[0][0][0]
        self.assertEqual(str(first_call_arg), "SHOW DATABASES;")

        # Check that the second call is CREATE DATABASE
        second_call_arg = call_args_list[1][0][0]
        self.assertEqual(str(second_call_arg), "CREATE DATABASE test_db ")

    @patch("nettacker.database.mysql.create_engine")
    def test_mysql_create_database_already_exists(self, mock_create_engine):
        """Test when database already exists"""
        # Set up mock config
        Config.db = MagicMock()
        Config.db.as_dict.return_value = {
            "username": "test_user",
            "password": "test_pass",
            "host": "localhost",
            "port": "3306",
            "name": "test_db",
        }
        Config.db.name = "test_db"

        # Set up mock connection and execution
        mock_conn = MagicMock()
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine
        mock_engine.connect.return_value.__enter__.return_value = mock_conn

        # Mock database query results - database already exists
        mock_conn.execute.return_value = [("mysql",), ("information_schema",), ("test_db",)]

        # Call the function
        mysql_create_database()

        # Assertions
        mock_create_engine.assert_called_once_with(
            "mysql+pymysql://test_user:test_pass@localhost:3306"
        )

        # Check that execute was called once with SHOW DATABASES
        self.assertEqual(mock_conn.execute.call_count, 1)
        call_arg = mock_conn.execute.call_args[0][0]
        self.assertEqual(str(call_arg), "SHOW DATABASES;")

    @patch("nettacker.database.mysql.create_engine")
    def test_mysql_create_database_exception(self, mock_create_engine):
        """Test exception handling in create database"""
        # Set up mock config
        Config.db = MagicMock()
        Config.db.as_dict.return_value = {
            "username": "test_user",
            "password": "test_pass",
            "host": "localhost",
            "port": "3306",
            "name": "test_db",
        }

        # Set up mock to raise exception
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine
        mock_engine.connect.side_effect = SQLAlchemyError("Connection error")

        # Call the function (should not raise exception)
        with patch("builtins.print") as mock_print:
            mysql_create_database()
            mock_print.assert_called_once()

    @patch("nettacker.database.mysql.create_engine")
    def test_mysql_create_tables(self, mock_create_engine):
        """Test table creation function"""
        # Set up mock config
        Config.db = MagicMock()
        Config.db.as_dict.return_value = {
            "username": "test_user",
            "password": "test_pass",
            "host": "localhost",
            "port": "3306",
            "name": "test_db",
        }

        # Set up mock engine
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine

        # Call the function
        with patch.object(Base.metadata, "create_all") as mock_create_all:
            mysql_create_tables()

            # Assertions
            mock_create_engine.assert_called_once_with(
                "mysql+pymysql://test_user:test_pass@localhost:3306/test_db"
            )
            mock_create_all.assert_called_once_with(mock_engine)