Improve tests for the ephemeral build manager.

test/test_buildman.py

@@ -1,12 +1,15 @@
 import unittest
 import etcd
+import os.path
+import time
 
-from trollius import coroutine, get_event_loop, From, Future
+from trollius import coroutine, get_event_loop, From, Future, sleep
 from mock import Mock
-from functools import partial
+from threading import Event
 
 from buildman.manager.executor import BuilderExecutor
-from buildman.manager.ephemeral import EphemeralBuilderManager, ETCD_BUILDER_PREFIX
+from buildman.manager.ephemeral import (EphemeralBuilderManager, ETCD_BUILDER_PREFIX,
+                                        ETCD_EXPIRE_RESULT)
 from buildman.server import BuildJobResult
 from buildman.component.buildcomponent import BuildComponent
 
@@ -14,10 +17,6 @@ from buildman.component.buildcomponent import BuildComponent
 BUILD_UUID = 'deadbeef-dead-beef-dead-deadbeefdead'
 
 
-import logging
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
-
 def async_test(f):
   def wrapper(*args, **kwargs):
     coro = coroutine(f)
@@ -29,11 +28,17 @@ def async_test(f):
 class TestEphemeral(unittest.TestCase):
   def __init__(self, *args, **kwargs):
     self.etcd_client_mock = None
+    self.etcd_wait_event = Event()
     self.test_executor = None
     super(TestEphemeral, self).__init__(*args, **kwargs)
 
   def _create_mock_etcd_client(self, *args, **kwargs):
+    def hang_until_event(*args, **kwargs):
+      time.sleep(.01)  # 10ms to simulate network latency
+      self.etcd_wait_event.wait()
+
     self.etcd_client_mock = Mock(spec=etcd.Client, name='etcd.Client')
+    self.etcd_client_mock.watch = Mock(side_effect=hang_until_event)
     return self.etcd_client_mock
 
   def _create_mock_executor(self, *args, **kwargs):
@@ -61,6 +66,7 @@ class TestEphemeral(unittest.TestCase):
 
     self.old_etcd_client_klass = EphemeralBuilderManager._etcd_client_klass
     EphemeralBuilderManager._etcd_client_klass = self._create_mock_etcd_client
+    self.etcd_wait_event.clear()
 
     self.register_component_callback = Mock()
     self.uniregister_component_callback = Mock()
@@ -77,7 +83,13 @@ class TestEphemeral(unittest.TestCase):
 
     self.manager.initialize({'EXECUTOR': 'test'})
 
+    self.mock_job_key = os.path.join(ETCD_BUILDER_PREFIX, BUILD_UUID)
+
   def tearDown(self):
+    self.etcd_wait_event.set()
+
+    self.manager.shutdown()
+
     del EphemeralBuilderManager._executors['test']
     EphemeralBuilderManager._etcd_client_klass = self.old_etcd_client_klass
 
@@ -95,15 +107,53 @@ class TestEphemeral(unittest.TestCase):
 
     self.assertTrue(is_scheduled)
 
-    job_key = ETCD_BUILDER_PREFIX + mock_job.job_details['build_uuid']
     self.etcd_client_mock.read.assert_called_once_with(ETCD_BUILDER_PREFIX, recursive=True)
-    self.assertEqual(len(self.test_executor.start_builder.call_args_list), 1)
-    self.assertEqual(self.etcd_client_mock.write.call_args_list[0][0][0], job_key)
-    self.assertEqual(self.etcd_client_mock.write.call_args_list[1][0][0], job_key)
+    self.assertEqual(self.test_executor.start_builder.call_count, 1)
+    self.assertEqual(self.etcd_client_mock.write.call_args_list[0][0][0], self.mock_job_key)
+    self.assertEqual(self.etcd_client_mock.write.call_args_list[1][0][0], self.mock_job_key)
 
-    self.assertEqual(len(self.register_component_callback.call_args_list), 1)
+    self.assertEqual(self.register_component_callback.call_count, 1)
 
     yield From(self.manager.job_completed(mock_job, BuildJobResult.COMPLETE, test_component))
 
-    self.assertEqual(len(self.test_executor.stop_builder.call_args_list), 1)
-    self.etcd_client_mock.delete.assert_called_once_with(job_key)
+    self.assertEqual(self.test_executor.stop_builder.call_count, 1)
+    self.etcd_client_mock.delete.assert_called_once_with(self.mock_job_key)
+
+  @async_test
+  def test_expiring_worker(self):
+    # Test that we are watching before anything else happens
+    self.etcd_client_mock.watch.assert_called_once_with(ETCD_BUILDER_PREFIX, recursive=True)
+
+    # Send a signal to the callback that a worker has expired
+    expired_result = Mock(sepc=etcd.EtcdResult)
+    expired_result.action = ETCD_EXPIRE_RESULT
+    expired_result.key = self.mock_job_key
+    expired_result._prev_node = Mock(spec=etcd.EtcdResult)
+    expired_result._prev_node.value = {'builder_id': '1234'}
+    expired_future = Future()
+    expired_future.set_result(expired_result)
+
+    self.manager._handle_key_expiration(expired_future)
+
+    yield From(sleep(.01))
+
+    self.test_executor.stop_builder.assert_called_once_with('1234')
+    self.assertEqual(self.test_executor.stop_builder.call_count, 1)
+
+    self.etcd_client_mock.delete.assert_called_once_with(self.mock_job_key)
+
+  @async_test
+  def test_change_worker(self):
+    # Send a signal to the callback that a worker key has been changed
+    set_result = Mock(sepc=etcd.EtcdResult)
+    set_result.action = 'set'
+    set_result.key = self.mock_job_key
+    set_future = Future()
+    set_future.set_result(set_result)
+
+    self.manager._handle_key_expiration(set_future)
+
+    yield From(sleep(.01))
+
+    self.assertEquals(self.test_executor.stop_builder.call_count, 0)
+