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")
