def _send_request_and_receive_response(self, server_uri, file_path): key = '{}_{}'.format(file_path, time.time_ns()) client = NotificationClient(server_uri=server_uri, default_namespace=SCHEDULER_NAMESPACE) event = BaseEvent(key=key, event_type=SchedulerInnerEventType.PARSE_DAG_REQUEST.value, value=file_path) client.send_event(event) watcher: ResponseWatcher = ResponseWatcher() client.start_listen_event(key=key, event_type=SchedulerInnerEventType.PARSE_DAG_RESPONSE.value, watcher=watcher) res: BaseEvent = watcher.get_result() self.assertEquals(event.key, res.key) self.assertEquals(event.value, file_path)
class EventSchedulerClient(object): def __init__(self, server_uri=None, namespace=None, ns_client=None): if ns_client is None: self.ns_client = NotificationClient(server_uri, namespace) else: self.ns_client = ns_client @staticmethod def generate_id(id): return '{}_{}'.format(id, time.time_ns()) def trigger_parse_dag(self) -> bool: id = self.generate_id('') watcher: ResponseWatcher = ResponseWatcher() handler: ThreadEventWatcherHandle \ = self.ns_client.start_listen_event(key=id, event_type=SchedulerInnerEventType.PARSE_DAG_RESPONSE.value, namespace=SCHEDULER_NAMESPACE, watcher=watcher) self.ns_client.send_event( BaseEvent( key=id, event_type=SchedulerInnerEventType.PARSE_DAG_REQUEST.value, value='')) result = watcher.get_result() handler.stop() return True def schedule_dag(self, dag_id) -> ExecutionContext: id = self.generate_id(dag_id) watcher: ResponseWatcher = ResponseWatcher() handler: ThreadEventWatcherHandle \ = self.ns_client.start_listen_event(key=id, event_type=SchedulerInnerEventType.RESPONSE.value, namespace=SCHEDULER_NAMESPACE, watcher=watcher) self.ns_client.send_event( RequestEvent(request_id=id, body=RunDagMessage(dag_id).to_json()).to_event()) result: ResponseEvent = ResponseEvent.from_base_event( watcher.get_result()) handler.stop() return ExecutionContext(dagrun_id=result.body) def stop_dag_run(self, dag_id, context: ExecutionContext) -> ExecutionContext: id = self.generate_id(str(dag_id) + str(context.dagrun_id)) watcher: ResponseWatcher = ResponseWatcher() handler: ThreadEventWatcherHandle \ = self.ns_client.start_listen_event(key=id, event_type=SchedulerInnerEventType.RESPONSE.value, namespace=SCHEDULER_NAMESPACE, watcher=watcher) self.ns_client.send_event( RequestEvent( request_id=id, body=StopDagRunMessage( dag_id=dag_id, dagrun_id=context.dagrun_id).to_json()).to_event()) result: ResponseEvent = ResponseEvent.from_base_event( watcher.get_result()) handler.stop() return ExecutionContext(dagrun_id=result.body) def schedule_task(self, dag_id: str, task_id: str, action: SchedulingAction, context: ExecutionContext) -> ExecutionContext: id = self.generate_id(context.dagrun_id) watcher: ResponseWatcher = ResponseWatcher() handler: ThreadEventWatcherHandle \ = self.ns_client.start_listen_event(key=id, event_type=SchedulerInnerEventType.RESPONSE.value, namespace=SCHEDULER_NAMESPACE, watcher=watcher) self.ns_client.send_event( RequestEvent(request_id=id, body=ExecuteTaskMessage( dag_id=dag_id, task_id=task_id, dagrun_id=context.dagrun_id, action=action.value).to_json()).to_event()) result: ResponseEvent = ResponseEvent.from_base_event( watcher.get_result()) handler.stop() return ExecutionContext(dagrun_id=result.body)
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))