class BaseExecutorTest(unittest.TestCase): def setUp(self): clear_db_event_model() self.master = NotificationMaster(service=NotificationService(EventModelStorage())) self.master.run() self.client = NotificationClient(server_uri="localhost:50051") def tearDown(self) -> None: self.master.stop() def test_get_event_buffer(self): executor = BaseExecutor() date = datetime.utcnow() try_number = 1 key1 = ("my_dag1", "my_task1", date, try_number) key2 = ("my_dag2", "my_task1", date, try_number) key3 = ("my_dag2", "my_task2", date, try_number) state = State.SUCCESS executor.event_buffer[key1] = state executor.event_buffer[key2] = state executor.event_buffer[key3] = state self.assertEqual(len(executor.get_event_buffer(("my_dag1",))), 1) self.assertEqual(len(executor.get_event_buffer()), 2) self.assertEqual(len(executor.event_buffer), 0) @mock.patch('airflow.executors.base_executor.BaseExecutor.sync') @mock.patch('airflow.executors.base_executor.BaseExecutor.trigger_tasks') @mock.patch('airflow.settings.Stats.gauge') def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock_sync): executor = BaseExecutor() executor.heartbeat() calls = [mock.call('executor.open_slots', mock.ANY), mock.call('executor.queued_tasks', mock.ANY), mock.call('executor.running_tasks', mock.ANY)] mock_stats_gauge.assert_has_calls(calls) def test_use_nf_executor(self): executor = BaseExecutor() executor.set_use_nf(True) executor.change_state('key', State.RUNNING) executor.change_state('key', State.SUCCESS) events = self.client.list_all_events(1) self.assertEqual(2, len(events))
class NotificationTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.storage = EventModelStorage() cls.master = NotificationMaster(NotificationService(cls.storage)) cls.master.run() @classmethod def tearDownClass(cls): cls.master.stop() def setUp(self): self.storage.clean_up() self.client = NotificationClient(server_uri="localhost:50051") def tearDown(self): self.client.stop_listen_events() self.client.stop_listen_event() def test_send_event(self): event = self.client.send_event(Event(key="key", value="value1")) self.assertTrue(event.version > 0) def test_list_events(self): event1 = self.client.send_event(Event(key="key", value="value1")) event2 = self.client.send_event(Event(key="key", value="value2")) event3 = self.client.send_event(Event(key="key", value="value3")) events = self.client.list_events("key", version=event1.version) self.assertEqual(2, len(events)) def test_listen_events(self): event_list = [] class TestWatch(EventWatcher): def __init__(self, event_list) -> None: super().__init__() self.event_list = event_list def process(self, events: List[Event]): self.event_list.extend(events) event1 = self.client.send_event(Event(key="key", value="value1")) self.client.start_listen_event(key="key", watcher=TestWatch(event_list), version=event1.version) event = self.client.send_event(Event(key="key", value="value2")) event = self.client.send_event(Event(key="key", value="value3")) self.client.stop_listen_event("key") events = self.client.list_events("key", version=event1.version) self.assertEqual(2, len(events)) self.assertEqual(2, len(event_list)) def test_all_listen_events(self): event = self.client.send_event(Event(key="key", value="value1")) event = self.client.send_event(Event(key="key", value="value2")) start_time = event.create_time event = self.client.send_event(Event(key="key", value="value3")) events = self.client.list_all_events(start_time) self.assertEqual(2, len(events)) def test_listen_all_events(self): event_list = [] class TestWatch(EventWatcher): def __init__(self, event_list) -> None: super().__init__() self.event_list = event_list def process(self, events: List[Event]): self.event_list.extend(events) try: self.client.start_listen_events(watcher=TestWatch(event_list)) event = self.client.send_event(Event(key="key1", value="value1")) event = self.client.send_event(Event(key="key2", value="value2")) event = self.client.send_event(Event(key="key3", value="value3")) finally: self.client.stop_listen_events() self.assertEqual(3, len(event_list))
class HaServerTest(unittest.TestCase): @classmethod def start_master(cls, host, port): port = str(port) server_uri = host + ":" + port storage = DbEventStorage() ha_manager = SimpleNotificationServerHaManager() ha_storage = DbHighAvailabilityStorage() service = HighAvailableNotificationService( storage, ha_manager, server_uri, ha_storage) master = NotificationMaster(service, port=int(port)) master.run() return master @classmethod def setUpClass(cls): cls.storage = DbEventStorage() cls.master1 = None cls.master2 = None cls.master3 = None def setUp(self): self.storage.clean_up() self.master1 = self.start_master("localhost", "50051") self.client = NotificationClient(server_uri="localhost:50051", enable_ha=True) def tearDown(self): self.client.stop_listen_events() self.client.stop_listen_event() self.client.disable_high_availability() if self.master1 is not None: self.master1.stop() if self.master2 is not None: self.master2.stop() if self.master3 is not None: self.master3.stop() def wait_for_new_members_detected(self, new_member_uri): for i in range(100): living_member = self.client.living_members if new_member_uri in living_member: break else: time.sleep(0.1) def test_server_change(self): self.client.send_event(BaseEvent(key="key", value="value1")) self.client.send_event(BaseEvent(key="key", value="value2")) self.client.send_event(BaseEvent(key="key", value="value3")) results = self.client.list_all_events() self.assertEqual(self.client.current_uri, "localhost:50051") self.master2 = self.start_master("localhost", "50052") self.wait_for_new_members_detected("localhost:50052") self.master1.stop() results2 = self.client.list_all_events() self.assertEqual(results, results2) self.assertEqual(self.client.current_uri, "localhost:50052") self.master3 = self.start_master("localhost", "50053") self.wait_for_new_members_detected("localhost:50053") self.master2.stop() results3 = self.client.list_all_events() self.assertEqual(results2, results3) self.assertEqual(self.client.current_uri, "localhost:50053") def test_send_listening_on_different_server(self): event_list = [] class TestWatch(EventWatcher): def __init__(self, event_list) -> None: super().__init__() self.event_list = event_list def process(self, events: List[BaseEvent]): self.event_list.extend(events) self.master2 = self.start_master("localhost", "50052") self.wait_for_new_members_detected("localhost:50052") another_client = NotificationClient(server_uri="localhost:50052") try: event1 = another_client.send_event(BaseEvent(key="key1", value="value1")) self.client.start_listen_events(watcher=TestWatch(event_list), version=event1.version) another_client.send_event(BaseEvent(key="key2", value="value2")) another_client.send_event(BaseEvent(key="key3", value="value3")) finally: self.client.stop_listen_events() self.assertEqual(2, len(event_list)) def test_start_with_multiple_servers(self): self.client.disable_high_availability() self.client = NotificationClient(server_uri="localhost:55001,localhost:50051", enable_ha=True) self.assertTrue(self.client.current_uri, "localhost:50051")