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
7import mock
8import unittest
9
10import common
11
12from autotest_lib.database import database_connection
13from autotest_lib.frontend import setup_django_environment
14from autotest_lib.frontend.afe import readonly_connection
15from autotest_lib.server import utils as server_utils
16from autotest_lib.scheduler import scheduler_lib
17from django.db import utils as django_utils
18
19
20class 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
94if __name__ == '__main__':
95    unittest.main()
96