X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/ffe3cdbc8c37e2b4a4e3ea4f67c1c9ca5d81e2ed..791bfa0fabd9b37ff8b52fb934f16b47461731da:/sdk/python/tests/test_websockets.py diff --git a/sdk/python/tests/test_websockets.py b/sdk/python/tests/test_websockets.py index 6b57fe3c8e..907dd93100 100644 --- a/sdk/python/tests/test_websockets.py +++ b/sdk/python/tests/test_websockets.py @@ -1,32 +1,227 @@ -import run_test_server -import unittest import arvados import arvados.events +from datetime import datetime, timedelta, tzinfo +import logging +import logging.handlers +import mock +import Queue +import run_test_server +import StringIO +import tempfile +import threading import time +import unittest + +class WebsocketTest(run_test_server.TestCaseWithServers): + MAIN_SERVER = {} + + TIME_PAST = time.time()-3600 + TIME_FUTURE = time.time()+3600 -class WebsocketTest(unittest.TestCase): def setUp(self): - run_test_server.run(websockets=True) - - def on_event(self, ev): - if self.state == 1: - self.assertEqual(200, ev['status']) - self.state = 2 - elif self.state == 2: - self.assertEqual(self.h[u'uuid'], ev[u'object_uuid']) - self.state = 3 - elif self.state == 3: - self.fail() - - def runTest(self): - self.state = 1 - - run_test_server.authorize_with("admin") - api = arvados.api('v1', cache=False) - arvados.events.subscribe(api, [['object_uuid', 'is_a', 'arvados#human']], lambda ev: self.on_event(ev)) - time.sleep(1) - self.h = api.humans().create(body={}).execute() - time.sleep(1) + self.ws = None def tearDown(self): - run_test_server.stop() + try: + if self.ws: + self.ws.close() + except Exception as e: + print("Error in teardown: ", e) + super(WebsocketTest, self).tearDown() + run_test_server.reset() + + def _test_subscribe(self, poll_fallback, expect_type, start_time=None, expected=1): + run_test_server.authorize_with('active') + events = Queue.Queue(100) + + # Create ancestor before subscribing. + # When listening with start_time in the past, this should also be retrieved. + # However, when start_time is omitted in subscribe, this should not be fetched. + ancestor = arvados.api('v1').humans().create(body={}).execute() + + filters = [['object_uuid', 'is_a', 'arvados#human']] + if start_time: + filters.append(['created_at', '>=', start_time]) + + self.ws = arvados.events.subscribe( + arvados.api('v1'), filters, + events.put_nowait, + poll_fallback=poll_fallback, + last_log_id=(1 if start_time else None)) + self.assertIsInstance(self.ws, expect_type) + self.assertEqual(200, events.get(True, 5)['status']) + human = arvados.api('v1').humans().create(body={}).execute() + + log_object_uuids = [] + for i in range(0, expected): + log_object_uuids.append(events.get(True, 5)['object_uuid']) + + if expected > 0: + self.assertIn(human['uuid'], log_object_uuids) + + if expected > 1: + self.assertIn(ancestor['uuid'], log_object_uuids) + + with self.assertRaises(Queue.Empty): + # assertEqual just serves to show us what unexpected thing + # comes out of the queue when the assertRaises fails; when + # the test passes, this assertEqual doesn't get called. + self.assertEqual(events.get(True, 2), None) + + def test_subscribe_websocket(self): + self._test_subscribe( + poll_fallback=False, expect_type=arvados.events.EventClient, expected=1) + + @mock.patch('arvados.events.EventClient.__init__') + def test_subscribe_poll(self, event_client_constr): + event_client_constr.side_effect = Exception('All is well') + self._test_subscribe( + poll_fallback=0.25, expect_type=arvados.events.PollClient, expected=1) + + def test_subscribe_websocket_with_start_time_past(self): + self._test_subscribe( + poll_fallback=False, expect_type=arvados.events.EventClient, + start_time=self.localiso(self.TIME_PAST), + expected=2) + + @mock.patch('arvados.events.EventClient.__init__') + def test_subscribe_poll_with_start_time_past(self, event_client_constr): + event_client_constr.side_effect = Exception('All is well') + self._test_subscribe( + poll_fallback=0.25, expect_type=arvados.events.PollClient, + start_time=self.localiso(self.TIME_PAST), + expected=2) + + def test_subscribe_websocket_with_start_time_future(self): + self._test_subscribe( + poll_fallback=False, expect_type=arvados.events.EventClient, + start_time=self.localiso(self.TIME_FUTURE), + expected=0) + + @mock.patch('arvados.events.EventClient.__init__') + def test_subscribe_poll_with_start_time_future(self, event_client_constr): + event_client_constr.side_effect = Exception('All is well') + self._test_subscribe( + poll_fallback=0.25, expect_type=arvados.events.PollClient, + start_time=self.localiso(self.TIME_FUTURE), + expected=0) + + def test_subscribe_websocket_with_start_time_past_utc(self): + self._test_subscribe( + poll_fallback=False, expect_type=arvados.events.EventClient, + start_time=self.utciso(self.TIME_PAST), + expected=2) + + def test_subscribe_websocket_with_start_time_future_utc(self): + self._test_subscribe( + poll_fallback=False, expect_type=arvados.events.EventClient, + start_time=self.utciso(self.TIME_FUTURE), + expected=0) + + def utciso(self, t): + return time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(t)) + + def localiso(self, t): + return time.strftime('%Y-%m-%dT%H:%M:%S', time.localtime(t)) + self.isotz(-time.timezone/60) + + def isotz(self, offset): + """Convert minutes-east-of-UTC to ISO8601 time zone designator""" + return '{:+03d}{:02d}'.format(offset/60, offset%60) + + # Test websocket reconnection on (un)execpted close + def _test_websocket_reconnect(self, close_unexpected): + run_test_server.authorize_with('active') + events = Queue.Queue(100) + + logstream = StringIO.StringIO() + rootLogger = logging.getLogger() + streamHandler = logging.StreamHandler(logstream) + rootLogger.addHandler(streamHandler) + + filters = [['object_uuid', 'is_a', 'arvados#human']] + filters.append(['created_at', '>=', self.localiso(self.TIME_PAST)]) + self.ws = arvados.events.subscribe( + arvados.api('v1'), filters, + events.put_nowait, + poll_fallback=False, + last_log_id=None) + self.assertIsInstance(self.ws, arvados.events.EventClient) + self.assertEqual(200, events.get(True, 5)['status']) + + # create obj + human = arvados.api('v1').humans().create(body={}).execute() + + # expect an event + self.assertIn(human['uuid'], events.get(True, 5)['object_uuid']) + with self.assertRaises(Queue.Empty): + self.assertEqual(events.get(True, 2), None) + + # close (im)properly + if close_unexpected: + self.ws.close_connection() + else: + self.ws.close() + + # create one more obj + human2 = arvados.api('v1').humans().create(body={}).execute() + + # (un)expect the object creation event + if close_unexpected: + log_object_uuids = [] + for i in range(0, 2): + event = events.get(True, 5) + if event.get('object_uuid') != None: + log_object_uuids.append(event['object_uuid']) + with self.assertRaises(Queue.Empty): + self.assertEqual(events.get(True, 2), None) + self.assertNotIn(human['uuid'], log_object_uuids) + self.assertIn(human2['uuid'], log_object_uuids) + else: + with self.assertRaises(Queue.Empty): + self.assertEqual(events.get(True, 2), None) + + # verify log message to ensure that an (un)expected close + log_messages = logstream.getvalue() + closeLogFound = log_messages.find("Unexpected close. Reconnecting.") + retryLogFound = log_messages.find("Error during websocket reconnect. Will retry") + if close_unexpected: + self.assertNotEqual(closeLogFound, -1) + else: + self.assertEqual(closeLogFound, -1) + rootLogger.removeHandler(streamHandler) + + def test_websocket_reconnect_on_unexpected_close(self): + self._test_websocket_reconnect(True) + + def test_websocket_no_reconnect_on_close_by_user(self): + self._test_websocket_reconnect(False) + + # Test websocket reconnection retry + @mock.patch('arvados.events._EventClient.connect') + def test_websocket_reconnect_retry(self, event_client_connect): + event_client_connect.side_effect = [None, Exception('EventClient.connect error'), None] + + logstream = StringIO.StringIO() + rootLogger = logging.getLogger() + streamHandler = logging.StreamHandler(logstream) + rootLogger.addHandler(streamHandler) + + run_test_server.authorize_with('active') + events = Queue.Queue(100) + + filters = [['object_uuid', 'is_a', 'arvados#human']] + self.ws = arvados.events.subscribe( + arvados.api('v1'), filters, + events.put_nowait, + poll_fallback=False, + last_log_id=None) + self.assertIsInstance(self.ws, arvados.events.EventClient) + + # simulate improper close + self.ws.on_closed() + + # verify log messages to ensure retry happened + log_messages = logstream.getvalue() + found = log_messages.find("Error 'EventClient.connect error' during websocket reconnect. Will retry") + self.assertNotEqual(found, -1) + rootLogger.removeHandler(streamHandler)