Home | History | Annotate | Download | only in scheduler 
      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