[autotest] Fixing and re-enabling monitor_db_functional_test.

The test was disabled and outdated. Database access and mocking of the drone
manager changed. This fixes these issues, updates the unit tests to the
current status and reanables them.

BUG=chromium:395756
DEPLOY=scheduler
TEST=ran ./utils/unittest_suite.py

Change-Id: I6a3eda5ddfaf07f06d6b403692b004b22939ffb6
Reviewed-on: https://chromium-review.googlesource.com/209567
Reviewed-by: Alex Miller <milleral@chromium.org>
Tested-by: Jakob Jülich <jakobjuelich@google.com>
Commit-Queue: Jakob Jülich <jakobjuelich@google.com>
diff --git a/scheduler/agent_task.py b/scheduler/agent_task.py
index 144f127..70bb6bd 100644
--- a/scheduler/agent_task.py
+++ b/scheduler/agent_task.py
@@ -121,8 +121,6 @@
 from autotest_lib.server import autoserv_utils
 
 
-_drone_manager = drone_manager.instance()
-
 AUTOSERV_NICE_LEVEL = 10
 
 
@@ -138,6 +136,7 @@
         """
         @param log_file_name: (optional) name of file to log command output to
         """
+        self._drone_manager = drone_manager.instance()
         self.done = False
         self.started = False
         self.success = None
@@ -394,13 +393,13 @@
 
 
     def register_necessary_pidfiles(self):
-        pidfile_id = _drone_manager.get_pidfile_id_from(
+        pidfile_id = self._drone_manager.get_pidfile_id_from(
                 self._working_directory(), self._pidfile_name())
-        _drone_manager.register_pidfile(pidfile_id)
+        self._drone_manager.register_pidfile(pidfile_id)
 
         paired_pidfile_id = self._paired_with_monitor().pidfile_id
         if paired_pidfile_id:
-            _drone_manager.register_pidfile(paired_pidfile_id)
+            self._drone_manager.register_pidfile(paired_pidfile_id)
 
 
     def recover(self):
@@ -464,7 +463,7 @@
         assert self.monitor
         if not self.monitor.has_process():
             return
-        _drone_manager.write_lines_to_file(
+        self._drone_manager.write_lines_to_file(
             self._keyval_path(), [self._format_keyval(field, value)],
             paired_with_process=self.monitor.get_process())
 
@@ -482,7 +481,7 @@
                                     for key, value in keyval_dict.iteritems())
         # always end with a newline to allow additional keyvals to be written
         keyval_contents += '\n'
-        _drone_manager.attach_file_to_execution(self._working_directory(),
+        self._drone_manager.attach_file_to_execution(self._working_directory(),
                                                 keyval_contents,
                                                 file_path=keyval_path)
 
@@ -587,10 +586,10 @@
                 source_path=self._working_directory() + '/',
                 destination_path=self.queue_entry.execution_path() + '/')
 
-        pidfile_id = _drone_manager.get_pidfile_id_from(
+        pidfile_id = self._drone_manager.get_pidfile_id_from(
                 self.queue_entry.execution_path(),
                 pidfile_name=drone_manager.AUTOSERV_PID_FILE)
-        _drone_manager.register_pidfile(pidfile_id)
+        self._drone_manager.register_pidfile(pidfile_id)
 
         if self.queue_entry.job.parse_failed_repair:
             self._parse_results([self.queue_entry])
@@ -615,7 +614,7 @@
             if self.monitor.has_process():
                 self._copy_results([self.task])
             if self.monitor.pidfile_id is not None:
-                _drone_manager.unregister_pidfile(self.monitor.pidfile_id)
+                self._drone_manager.unregister_pidfile(self.monitor.pidfile_id)
 
 
     def remove_special_tasks(self, special_task_to_remove, keep_last_one=False):
diff --git a/scheduler/drone_manager_unittest.py b/scheduler/drone_manager_unittest.py
index f06e15a..e9f1f53 100755
--- a/scheduler/drone_manager_unittest.py
+++ b/scheduler/drone_manager_unittest.py
@@ -443,7 +443,7 @@
         # Our manager instance isn't the drone manager singletone that the
         # pidfile_monitor will use by default, becuase setUp doesn't call
         # drone_manager.instance().
-        self.god.stub_with(pidfile_monitor, '_drone_manager', self.manager)
+        self.god.stub_with(drone_manager, '_the_instance', self.manager)
         monitor = pidfile_monitor.PidfileRunMonitor()
         monitor.pidfile_id = drone_manager.PidfileId(pidfile_path)
         self.manager.register_pidfile(monitor.pidfile_id)
diff --git a/scheduler/monitor_db.py b/scheduler/monitor_db.py
index d87262d..e785c04 100755
--- a/scheduler/monitor_db.py
+++ b/scheduler/monitor_db.py
@@ -51,6 +51,7 @@
 """
 
 _db_manager = None
+_db = None
 _shutdown = False
 
 # These 2 globals are replaced for testing
@@ -192,6 +193,8 @@
     os.environ['PATH'] = AUTOTEST_SERVER_DIR + ':' + os.environ['PATH']
     global _db_manager
     _db_manager = scheduler_lib.ConnectionManager()
+    global _db
+    _db = _db_manager.get_connection()
     logging.info("Setting signal handler")
     signal.signal(signal.SIGINT, handle_signal)
     signal.signal(signal.SIGTERM, handle_signal)
@@ -241,9 +244,9 @@
         self._last_clean_time = time.time()
         user_cleanup_time = scheduler_config.config.clean_interval_minutes
         self._periodic_cleanup = monitor_db_cleanup.UserCleanup(
-                _db_manager.get_connection(), user_cleanup_time)
+                _db, user_cleanup_time)
         self._24hr_upkeep = monitor_db_cleanup.TwentyFourHourUpkeep(
-                _db_manager.get_connection())
+                _db)
         self._host_agents = {}
         self._queue_entry_agents = {}
         self._tick_count = 0
diff --git a/scheduler/monitor_db_functional_test.py b/scheduler/monitor_db_functional_test.py
index ab341f0..a53af24 100755
--- a/scheduler/monitor_db_functional_test.py
+++ b/scheduler/monitor_db_functional_test.py
@@ -1,6 +1,6 @@
 #!/usr/bin/python
 
-import logging, os, unittest
+import logging, os, signal, unittest
 import common
 from autotest_lib.client.common_lib import enum, global_config, host_protections
 from autotest_lib.database import database_connection
@@ -64,6 +64,19 @@
                                 in _PIDFILE_TO_PIDFILE_TYPE.iteritems())
 
 
+class MockConnectionManager(object):
+    """docstring for MockConnectionManager"""
+
+    db = None
+
+    def __init__(self):
+        super(MockConnectionManager, self).__init__()
+
+    def get_connection(self):
+        assert MockConnectionManager.db
+        return MockConnectionManager.db
+
+
 class MockDroneManager(NullMethodObject):
     """
     Public attributes:
@@ -71,7 +84,7 @@
             tests can change this to activate throttling.
     """
     _NULL_METHODS = ('reinitialize_drones', 'copy_to_results_repository',
-                     'copy_results_on_drone')
+                     'copy_results_on_drone', 'trigger_refresh', 'sync_refresh')
 
     class _DummyPidfileId(object):
         """
@@ -115,11 +128,12 @@
         # maps process to pidfile IDs
         self._process_index = {}
         # tracks pidfiles of processes that have been killed
-        self._killed_pidfiles = set()
+        self._pids_to_signals_received = {}
         # pidfile IDs that have just been unregistered (so will disappear on the
         # next cycle)
         self._unregistered_pidfiles = set()
-
+        # Pids to write exit status for at end of tick
+        self._set_pidfile_exit_status_queue = []
 
     # utility APIs for use by the test
 
@@ -140,9 +154,9 @@
         contents.num_tests_failed = 0
 
 
-    def was_last_process_killed(self, pidfile_type):
+    def was_last_process_killed(self, pidfile_type, sigs):
         pidfile_id = self._last_pidfile_id[pidfile_type]
-        return pidfile_id in self._killed_pidfiles
+        return sigs == self._pids_to_signals_received[pidfile_id]
 
 
     def nonfinished_pidfile_ids(self):
@@ -201,6 +215,10 @@
             self._process_index[process] = pidfile_id
         self._future_pidfiles = []
 
+        for pidfile_id in self._set_pidfile_exit_status_queue:
+            self._set_pidfile_exit_status(pidfile_id, 271)
+        self._set_pidfile_exit_status_queue = []
+
 
     def attach_file_to_execution(self, result_dir, file_contents,
                                  file_path=None):
@@ -281,10 +299,15 @@
                                        default_pidfile)
 
 
-    def kill_process(self, process):
+    def kill_process(self, process, sig=signal.SIGKILL):
         pidfile_id = self._process_index[process]
-        self._killed_pidfiles.add(pidfile_id)
-        self._set_pidfile_exit_status(pidfile_id, 271)
+
+        if pidfile_id not in self._pids_to_signals_received:
+            self._pids_to_signals_received[pidfile_id] = set()
+        self._pids_to_signals_received[pidfile_id].add(sig)
+
+        if signal.SIGKILL == sig:
+            self._set_pidfile_exit_status_queue.append(pidfile_id)
 
 
 class MockEmailManager(NullMethodObject):
@@ -336,6 +359,9 @@
         self.god.stub_with(monitor_db, '_db', self._database)
         self.god.stub_with(scheduler_models, '_db', self._database)
 
+        MockConnectionManager.db = self._database
+        scheduler_lib.ConnectionManager = MockConnectionManager
+
         monitor_db.initialize_globals()
         scheduler_models.initialize_globals()
 
@@ -540,6 +566,7 @@
 
 
     def _finish_job(self, queue_entry):
+        self._check_statuses(queue_entry, HqeStatus.RUNNING)
         self.mock_drone_manager.finish_process(_PidfileType.JOB)
         self._run_dispatcher() # launches parsing
         self._check_statuses(queue_entry, HqeStatus.PARSING)
@@ -650,11 +677,13 @@
     def test_job_abort_in_verify(self):
         self._initialize_test()
         job = self._create_job(hosts=[1])
+        queue_entries = list(job.hostqueueentry_set.all())
         self._run_dispatcher() # launches verify
+        self._check_statuses(queue_entries[0], HqeStatus.VERIFYING)
         job.hostqueueentry_set.update(aborted=True)
         self._run_dispatcher() # kills verify, launches cleanup
         self.assert_(self.mock_drone_manager.was_last_process_killed(
-                _PidfileType.VERIFY))
+                _PidfileType.VERIFY, set([signal.SIGKILL])))
         self.mock_drone_manager.finish_process(_PidfileType.CLEANUP)
         self._run_dispatcher()
 
@@ -664,12 +693,17 @@
         job = self._create_job(hosts=[1])
         job.run_reset = False
         job.save()
+        queue_entries = list(job.hostqueueentry_set.all())
 
         self._run_dispatcher() # launches job
+
+        self._check_statuses(queue_entries[0], HqeStatus.RUNNING)
+
         job.hostqueueentry_set.update(aborted=True)
+
         self._run_dispatcher() # kills job, launches gathering
-        self.assert_(self.mock_drone_manager.was_last_process_killed(
-                _PidfileType.JOB))
+
+        self._check_statuses(queue_entries[0], HqeStatus.GATHERING)
         self.mock_drone_manager.finish_process(_PidfileType.GATHER)
         self._run_dispatcher() # launches parsing + cleanup
         queue_entry = job.hostqueueentry_set.all()[0]
@@ -694,13 +728,17 @@
 
     def test_no_pidfile_leaking(self):
         self._initialize_test()
+
         self.test_simple_job()
+        self.mock_drone_manager.refresh()
         self.assertEquals(self.mock_drone_manager._pidfiles, {})
 
         self.test_job_abort_in_verify()
+        self.mock_drone_manager.refresh()
         self.assertEquals(self.mock_drone_manager._pidfiles, {})
 
         self.test_job_abort()
+        self.mock_drone_manager.refresh()
         self.assertEquals(self.mock_drone_manager._pidfiles, {})
 
 
@@ -936,16 +974,17 @@
 
         self.mock_drone_manager.process_capacity = 2
         self._run_dispatcher() # verify runs on 1 and 2
-        _check_hqe_statuses(HqeStatus.VERIFYING, HqeStatus.VERIFYING,
-                            HqeStatus.QUEUED)
+        queue_entries = list(job.hostqueueentry_set.all())
+        _check_hqe_statuses(HqeStatus.QUEUED,
+                            HqeStatus.VERIFYING, HqeStatus.VERIFYING)
         self.assertEquals(len(self.mock_drone_manager.running_pidfile_ids()), 2)
 
         self.mock_drone_manager.finish_specific_process(
-                'hosts/host1/1-verify', drone_manager.AUTOSERV_PID_FILE)
+                'hosts/host3/1-verify', drone_manager.AUTOSERV_PID_FILE)
         self.mock_drone_manager.finish_process(_PidfileType.VERIFY)
         self._run_dispatcher() # verify runs on 3
-        _check_hqe_statuses(HqeStatus.PENDING, HqeStatus.PENDING,
-                            HqeStatus.VERIFYING)
+        _check_hqe_statuses(HqeStatus.VERIFYING, HqeStatus.PENDING,
+                            HqeStatus.PENDING)
 
         self.mock_drone_manager.finish_process(_PidfileType.VERIFY)
         self._run_dispatcher() # job won't run due to throttling
@@ -1001,26 +1040,6 @@
                              HostStatus.CLEANING)
 
 
-    def test_simple_atomic_group_job(self):
-        job = self._create_job(atomic_group=1)
-        self._run_dispatcher() # expand + verify
-        queue_entries = job.hostqueueentry_set.all()
-        self.assertEquals(len(queue_entries), 2)
-        self.assertEquals(queue_entries[0].host.hostname, 'host5')
-        self.assertEquals(queue_entries[1].host.hostname, 'host6')
-
-        self.mock_drone_manager.finish_process(_PidfileType.VERIFY)
-        self._run_dispatcher() # delay task started waiting
-
-        self.mock_drone_manager.finish_specific_process(
-                'hosts/host6/1-verify', drone_manager.AUTOSERV_PID_FILE)
-        self._run_dispatcher() # job starts now
-        for entry in queue_entries:
-            self._check_statuses(entry, HqeStatus.RUNNING, HostStatus.RUNNING)
-
-        # rest of job proceeds normally
-
-
     def test_simple_metahost_assignment(self):
         job = self._create_job(metahosts=[1])
         self._run_dispatcher()
diff --git a/scheduler/monitor_db_unittest.py b/scheduler/monitor_db_unittest.py
index 184619e..2ead35b 100755
--- a/scheduler/monitor_db_unittest.py
+++ b/scheduler/monitor_db_unittest.py
@@ -536,7 +536,7 @@
         self.god = mock.mock_god()
         self.mock_drone_manager = self.god.create_mock_class(
             drone_manager.DroneManager, 'drone_manager')
-        self.god.stub_with(pidfile_monitor, '_drone_manager',
+        self.god.stub_with(drone_manager, '_the_instance',
                            self.mock_drone_manager)
         self.god.stub_function(email_manager.manager, 'enqueue_notify_email')
         self.god.stub_with(pidfile_monitor, '_get_pidfile_timeout_secs',
diff --git a/scheduler/pidfile_monitor.py b/scheduler/pidfile_monitor.py
index cce8792..c5e6395 100644
--- a/scheduler/pidfile_monitor.py
+++ b/scheduler/pidfile_monitor.py
@@ -11,7 +11,6 @@
 from autotest_lib.scheduler import scheduler_config
 
 
-_drone_manager = drone_manager.instance()
 
 def _get_pidfile_timeout_secs():
     """@returns How long to wait for autoserv to write pidfile."""
@@ -34,6 +33,7 @@
 
 
     def __init__(self):
+        self._drone_manager = drone_manager.instance()
         self.lost_process = False
         self._start_time = None
         self.pidfile_id = None
@@ -58,7 +58,7 @@
         if nice_level is not None:
             command = ['nice', '-n', str(nice_level)] + command
         self._set_start_time()
-        self.pidfile_id = _drone_manager.execute_command(
+        self.pidfile_id = self._drone_manager.execute_command(
             command, working_directory, pidfile_name=pidfile_name,
             num_processes=num_processes, log_file=log_file,
             paired_with_pidfile=paired_with_pidfile, username=username,
@@ -69,15 +69,15 @@
                                    pidfile_name=drone_manager.AUTOSERV_PID_FILE,
                                    num_processes=None):
         self._set_start_time()
-        self.pidfile_id = _drone_manager.get_pidfile_id_from(
+        self.pidfile_id = self._drone_manager.get_pidfile_id_from(
             execution_path, pidfile_name=pidfile_name)
         if num_processes is not None:
-            _drone_manager.declare_process_count(self.pidfile_id, num_processes)
+            self._drone_manager.declare_process_count(self.pidfile_id, num_processes)
 
 
     def kill(self):
         if self.has_process():
-            _drone_manager.kill_process(self.get_process())
+            self._drone_manager.kill_process(self.get_process())
             self._killed = True
 
 
@@ -95,7 +95,7 @@
     def _read_pidfile(self, use_second_read=False):
         assert self.pidfile_id is not None, (
             'You must call run() or attach_to_existing_process()')
-        contents = _drone_manager.get_pidfile_contents(
+        contents = self._drone_manager.get_pidfile_contents(
             self.pidfile_id, use_second_read=use_second_read)
         if contents.is_invalid():
             self._state = drone_manager.PidfileContents()
@@ -122,7 +122,7 @@
 
         if self._state.exit_status is None:
             # double check whether or not autoserv is running
-            if _drone_manager.is_process_running(self._state.process):
+            if self._drone_manager.is_process_running(self._state.process):
                 return
 
             # pid but no running process - maybe process *just* exited
@@ -195,11 +195,11 @@
     def try_copy_results_on_drone(self, **kwargs):
         if self.has_process():
             # copy results logs into the normal place for job results
-            _drone_manager.copy_results_on_drone(self.get_process(), **kwargs)
+            self._drone_manager.copy_results_on_drone(self.get_process(), **kwargs)
 
 
     def try_copy_to_results_repository(self, source, **kwargs):
         if self.has_process():
-            _drone_manager.copy_to_results_repository(self.get_process(),
+            self._drone_manager.copy_to_results_repository(self.get_process(),
                                                       source, **kwargs)
 
diff --git a/scheduler/postjob_task.py b/scheduler/postjob_task.py
index a0dd902..e38a06b 100644
--- a/scheduler/postjob_task.py
+++ b/scheduler/postjob_task.py
@@ -18,7 +18,6 @@
 from autotest_lib.server import autoserv_utils
 
 
-_drone_manager = drone_manager.instance()
 _parser_path = autoserv_utils._parser_path_func(
                 autoserv_utils.AUTOTEST_INSTALL_DIR)
 
@@ -37,7 +36,7 @@
     def _command_line(self):
         # Do we need testing_mode?
         return self._generate_command(
-                _drone_manager.absolute_path(self._working_directory()))
+                self._drone_manager.absolute_path(self._working_directory()))
 
 
     def _generate_command(self, results_dir):
@@ -383,7 +382,7 @@
             failed_file = os.path.join(self._working_directory(),
                                        self._ARCHIVING_FAILED_FILE)
             paired_process = self._paired_with_monitor().get_process()
-            _drone_manager.write_lines_to_file(
+            self._drone_manager.write_lines_to_file(
                     failed_file, ['Archiving failed with exit code %s'
                                   % self.monitor.exit_code()],
                     paired_with_process=paired_process)
diff --git a/utils/unittest_suite.py b/utils/unittest_suite.py
index ee9b3c8..f9c53a2 100755
--- a/utils/unittest_suite.py
+++ b/utils/unittest_suite.py
@@ -92,7 +92,6 @@
     'ap_configurator_test.py',
     'chaos_base_test.py',
     'chaos_interop_test.py',
-    'monitor_db_functional_test.py',
     'atomic_group_unittests.py',
     # crbug.com/251395
     'dev_server_test.py',