scheduler/scheduler_lib_unittest.py

#!/usr/bin/python
#
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import mock
import unittest

import common

from autotest_lib.database import database_connection
from autotest_lib.frontend import setup_django_environment
from autotest_lib.frontend.afe import readonly_connection
from autotest_lib.server import utils as server_utils
from autotest_lib.scheduler import scheduler_lib
from django.db import utils as django_utils


class ConnectionManagerTests(unittest.TestCase):
    """Connection manager unittests."""

    def setUp(self):
        self.connection_manager = None
        readonly_connection.set_globally_disabled = mock.MagicMock()
        setup_django_environment.enable_autocommit = mock.MagicMock()
        server_utils.Singleton._instances = {}


    def tearDown(self):
        readonly_connection.set_globally_disabled.reset_mock()
        setup_django_environment.enable_autocommit.reset_mock()


    def testConnectionDisconnect(self):
        """Test connection and disconnecting from the database."""
        # Test that the connection manager only opens a connection once.
        connection_manager = scheduler_lib.ConnectionManager()
        connection_manager.open_connection = mock.MagicMock()
        connection = connection_manager.get_connection()
        connection_manager.open_connection.assert_called_once_with()
        connection_manager.open_connection.reset_mock()
        connection = connection_manager.get_connection()
        self.assertTrue(
                connection_manager.open_connection.call_count == 0)
        connection_manager.open_connection.reset_mock()

        # Test that del on the connection manager closes the connection
        connection_manager.disconnect = mock.MagicMock()
        connection_manager.__del__()
        connection_manager.disconnect.assert_called_once_with()


    def testConnectionReconnect(self):
        """Test that retries don't destroy the connection."""
        database_connection._DjangoBackend.execute = mock.MagicMock()
        database_connection._DjangoBackend.execute.side_effect = (
                django_utils.DatabaseError('Database Error'))
        connection_manager = scheduler_lib.ConnectionManager()
        connection = connection_manager.get_connection()
        self.assertRaises(django_utils.DatabaseError,
                          connection.execute, *('', None, True))
        self.assertTrue(
                database_connection._DjangoBackend.execute.call_count == 2)
        database_connection._DjangoBackend.execute.reset_mock()
        self.assertTrue(connection_manager.db_connection ==
                        connection_manager.get_connection())


    def testConnectionManagerSingleton(self):
        """Test that the singleton works as expected."""
        # Confirm that instantiating the class applies global db settings.
        connection_manager = scheduler_lib.ConnectionManager()
        readonly_connection.set_globally_disabled.assert_called_once_with(True)
        setup_django_environment.enable_autocommit.assert_called_once_with()

        readonly_connection.set_globally_disabled.reset_mock()
        setup_django_environment.enable_autocommit.reset_mock()

        # Confirm that instantiating another connection manager doesn't change
        # the database settings, and in fact, returns the original manager.
        connection_manager_2 = scheduler_lib.ConnectionManager()
        self.assertTrue(connection_manager == connection_manager_2)
        self.assertTrue(
                readonly_connection.set_globally_disabled.call_count == 0)
        self.assertTrue(
                setup_django_environment.enable_autocommit.call_count == 0)

        # Confirm that we don't open the connection when the class is
        # instantiated.
        self.assertTrue(connection_manager.db_connection is None)


if __name__ == '__main__':
    unittest.main()