class HaDbStorageTest(unittest.TestCase, NotificationTest):
"""
This test is used to ensure the high availability would not break the original functionality.
"""
@classmethod
def set_up_class(cls):
cls.storage = DbEventStorage()
cls.master1 = start_ha_master("localhost", 50051)
# The server startup is asynchronous, we need to wait for a while
# to ensure it writes its metadata to the db.
time.sleep(0.1)
cls.master2 = start_ha_master("localhost", 50052)
time.sleep(0.1)
cls.master3 = start_ha_master("localhost", 50053)
time.sleep(0.1)
@classmethod
def setUpClass(cls):
cls.set_up_class()
@classmethod
def tearDownClass(cls):
cls.master1.stop()
cls.master2.stop()
cls.master3.stop()
def setUp(self):
self.storage.clean_up()
self.client = NotificationClient(server_uri="localhost:50052",
enable_ha=True,
list_member_interval_ms=1000,
retry_timeout_ms=10000)
def tearDown(self):
self.client.stop_listen_events()
self.client.stop_listen_event()
self.client.disable_high_availability()
def disable_high_availability(self):
if hasattr(self, "aiflow_ha_running"):
self.aiflow_ha_running = False
NotificationClient.disable_high_availability(self)
if hasattr(self, "aiflow_ha_running"):
self.list_aiflow_member_thread.join()
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")
