1#!/usr/bin/python
2
3import common
4import logging, unittest
5from autotest_lib.frontend import setup_django_environment
6from autotest_lib.database import database_connection
7from autotest_lib.frontend.afe import frontend_test_utils, models
8from autotest_lib.scheduler import monitor_db_cleanup, scheduler_config
9from autotest_lib.client.common_lib import host_protections
10
11class UserCleanupTest(unittest.TestCase, frontend_test_utils.FrontendTestMixin):
12    def setUp(self):
13        logging.basicConfig(level=logging.DEBUG)
14        self._frontend_common_setup()
15        self._database = (
16            database_connection.DatabaseConnection.get_test_database())
17        self._database.connect(db_type='django')
18        self.cleanup = monitor_db_cleanup.UserCleanup(self._database, 1)
19
20
21    def tearDown(self):
22        self._frontend_common_teardown()
23
24
25    def test_reverify_dead_hosts(self):
26        # unlimited reverifies
27        self.god.stub_with(scheduler_config.config,
28                           'reverify_max_hosts_at_once', 0)
29        for i in (0, 1, 2):
30            self.hosts[i].status = models.Host.Status.REPAIR_FAILED
31            self.hosts[i].save()
32
33        self.hosts[1].locked = True
34        self.hosts[1].save()
35
36        self.hosts[2].protection = host_protections.Protection.DO_NOT_VERIFY
37        self.hosts[2].save()
38
39        self.god.stub_with(self.cleanup, '_should_reverify_hosts_now',
40                           lambda : True)
41        self.cleanup._reverify_dead_hosts()
42
43        tasks = models.SpecialTask.objects.all()
44        self.assertEquals(len(tasks), 1)
45        self.assertEquals(tasks[0].host.id, 1)
46        self.assertEquals(tasks[0].task, models.SpecialTask.Task.VERIFY)
47
48
49    def test_reverify_dead_hosts_limits(self):
50        # limit the number of reverifies
51        self.assertTrue(hasattr(scheduler_config.config,
52                                'reverify_max_hosts_at_once'))
53        self.god.stub_with(scheduler_config.config,
54                           'reverify_max_hosts_at_once', 2)
55        for i in (0, 1, 2, 3, 4, 5):
56            self.hosts[i].status = models.Host.Status.REPAIR_FAILED
57            self.hosts[i].save()
58
59        self.hosts[1].locked = True
60        self.hosts[1].save()
61
62        self.hosts[2].protection = host_protections.Protection.DO_NOT_VERIFY
63        self.hosts[2].save()
64
65        self.god.stub_with(self.cleanup, '_should_reverify_hosts_now',
66                           lambda : True)
67        self.cleanup._reverify_dead_hosts()
68
69        tasks = models.SpecialTask.objects.all()
70        # four hosts need reverifying but our max limit was set to 2
71        self.assertEquals(len(tasks), 2)
72        self.assertTrue(tasks[0].host.id in (1, 4, 5, 6))
73        self.assertTrue(tasks[1].host.id in (1, 4, 5, 6))
74        self.assertEquals(tasks[0].task, models.SpecialTask.Task.VERIFY)
75        self.assertEquals(tasks[1].task, models.SpecialTask.Task.VERIFY)
76
77
78if __name__ == '__main__':
79    unittest.main()
80