| # 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.readonly_connection import ReadOnlyConnection |
| 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 |
| ReadOnlyConnection.set_globally_disabled = mock.MagicMock() |
| setup_django_environment.enable_autocommit = mock.MagicMock() |
| scheduler_lib.Singleton._instances = {} |
| |
| |
| def tearDown(self): |
| ReadOnlyConnection.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() |
| ReadOnlyConnection.set_globally_disabled.assert_called_once_with(False) |
| setup_django_environment.enable_autocommit.assert_called_once_with() |
| |
| ReadOnlyConnection.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( |
| ReadOnlyConnection.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) |