1 #!/usr/bin/python 2 # 3 # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 4 # Use of this source code is governed by a BSD-style license that can be 5 # found in the LICENSE file. 6 7 import mock 8 import unittest 9 10 import common 11 12 from autotest_lib.database import database_connection 13 from autotest_lib.frontend import setup_django_environment 14 from autotest_lib.frontend.afe import readonly_connection 15 from autotest_lib.server import utils as server_utils 16 from autotest_lib.scheduler import scheduler_lib 17 from django.db import utils as django_utils 18 19 20 class ConnectionManagerTests(unittest.TestCase): 21 """Connection manager unittests.""" 22 23 def setUp(self): 24 self.connection_manager = None 25 readonly_connection.set_globally_disabled = mock.MagicMock() 26 setup_django_environment.enable_autocommit = mock.MagicMock() 27 server_utils.Singleton._instances = {} 28 29 30 def tearDown(self): 31 readonly_connection.set_globally_disabled.reset_mock() 32 setup_django_environment.enable_autocommit.reset_mock() 33 34 35 def testConnectionDisconnect(self): 36 """Test connection and disconnecting from the database.""" 37 # Test that the connection manager only opens a connection once. 38 connection_manager = scheduler_lib.ConnectionManager() 39 connection_manager.open_connection = mock.MagicMock() 40 connection = connection_manager.get_connection() 41 connection_manager.open_connection.assert_called_once_with() 42 connection_manager.open_connection.reset_mock() 43 connection = connection_manager.get_connection() 44 self.assertTrue( 45 connection_manager.open_connection.call_count == 0) 46 connection_manager.open_connection.reset_mock() 47 48 # Test that del on the connection manager closes the connection 49 connection_manager.disconnect = mock.MagicMock() 50 connection_manager.__del__() 51 connection_manager.disconnect.assert_called_once_with() 52 53 54 def testConnectionReconnect(self): 55 """Test that retries don't destroy the connection.""" 56 database_connection._DjangoBackend.execute = mock.MagicMock() 57 database_connection._DjangoBackend.execute.side_effect = ( 58 django_utils.DatabaseError('Database Error')) 59 connection_manager = scheduler_lib.ConnectionManager() 60 connection = connection_manager.get_connection() 61 self.assertRaises(django_utils.DatabaseError, 62 connection.execute, *('', None, True)) 63 self.assertTrue( 64 database_connection._DjangoBackend.execute.call_count == 2) 65 database_connection._DjangoBackend.execute.reset_mock() 66 self.assertTrue(connection_manager.db_connection == 67 connection_manager.get_connection()) 68 69 70 def testConnectionManagerSingleton(self): 71 """Test that the singleton works as expected.""" 72 # Confirm that instantiating the class applies global db settings. 73 connection_manager = scheduler_lib.ConnectionManager() 74 readonly_connection.set_globally_disabled.assert_called_once_with(True) 75 setup_django_environment.enable_autocommit.assert_called_once_with() 76 77 readonly_connection.set_globally_disabled.reset_mock() 78 setup_django_environment.enable_autocommit.reset_mock() 79 80 # Confirm that instantiating another connection manager doesn't change 81 # the database settings, and in fact, returns the original manager. 82 connection_manager_2 = scheduler_lib.ConnectionManager() 83 self.assertTrue(connection_manager == connection_manager_2) 84 self.assertTrue( 85 readonly_connection.set_globally_disabled.call_count == 0) 86 self.assertTrue( 87 setup_django_environment.enable_autocommit.call_count == 0) 88 89 # Confirm that we don't open the connection when the class is 90 # instantiated. 91 self.assertTrue(connection_manager.db_connection is None) 92 93 94 if __name__ == '__main__': 95 unittest.main() 96