sdk/python/tests/test_websockets.py

   1 import Queue
   2 import run_test_server
   3 import unittest
   4 import arvados
   5 import arvados.events
   6 import mock
   7 import threading
   8 from datetime import datetime, timedelta
   9 import time
  10
  11 class WebsocketTest(run_test_server.TestCaseWithServers):
  12     MAIN_SERVER = {}
  13
  14     def setUp(self):
  15         self.ws = None
  16
  17     def tearDown(self):
  18         if self.ws:
  19             self.ws.close()
  20         super(WebsocketTest, self).tearDown()
  21
  22     def _test_subscribe(self, poll_fallback, expect_type, last_log_id=None, additional_filters=None, expected=1):
  23         run_test_server.authorize_with('active')
  24         events = Queue.Queue(100)
  25
  26         # Create ancestor before subscribing.
  27         # When listening with start_time in the past, this should also be retrieved.
  28         # However, when start_time is omitted in subscribe, this should not be fetched.
  29         ancestor = arvados.api('v1').humans().create(body={}).execute()
  30         time.sleep(5)
  31
  32         filters = [['object_uuid', 'is_a', 'arvados#human']]
  33         if additional_filters:
  34             filters = filters + additional_filters
  35
  36         self.ws = arvados.events.subscribe(
  37             arvados.api('v1'), filters,
  38             events.put, poll_fallback=poll_fallback, last_log_id=last_log_id)
  39         self.assertIsInstance(self.ws, expect_type)
  40         self.assertEqual(200, events.get(True, 5)['status'])
  41         human = arvados.api('v1').humans().create(body={}).execute()
  42
  43         if last_log_id == None or expected == 0:
  44             self.assertEqual(human['uuid'], events.get(True, 5)['object_uuid'])
  45             self.assertTrue(events.empty(), "got more events than expected")
  46         else:
  47             log_events = []
  48             for i in range(0, 20):
  49                 try:
  50                     event = events.get(True, 5)
  51                     self.assertTrue(event['object_uuid'] is not None)
  52                     log_events.append(event['object_uuid'])
  53                 except:
  54                     break;
  55
  56             self.assertTrue(len(log_events)>1)
  57             self.assertTrue(human['uuid'] in log_events)
  58             self.assertTrue(ancestor['uuid'] in log_events)
  59
  60     def test_subscribe_websocket(self):
  61         self._test_subscribe(
  62             poll_fallback=False, expect_type=arvados.events.EventClient)
  63
  64     @mock.patch('arvados.events.EventClient.__init__')
  65     def test_subscribe_poll(self, event_client_constr):
  66         event_client_constr.side_effect = Exception('All is well')
  67         self._test_subscribe(
  68             poll_fallback=1, expect_type=arvados.events.PollClient)
  69
  70     def test_subscribe_websocket_with_start_time_date_only(self):
  71         lastHour = datetime.today() - timedelta(hours = 1)
  72         self._test_subscribe(
  73             poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
  74                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d')]])
  75
  76     @mock.patch('arvados.events.EventClient.__init__')
  77     def test_poll_with_start_time_date_only(self, event_client_constr):
  78         event_client_constr.side_effect = Exception('All is well')
  79         lastHour = datetime.today() - timedelta(hours = 1)
  80         self._test_subscribe(
  81             poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
  82                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d')]])
  83
  84     def test_subscribe_websocket_with_start_time_last_hour(self):
  85         lastHour = datetime.today() - timedelta(hours = 1)
  86         self._test_subscribe(
  87             poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
  88                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d %H:%M:%S')]])
  89
  90     @mock.patch('arvados.events.EventClient.__init__')
  91     def test_subscribe_poll_with_start_time_last_hour(self, event_client_constr):
  92         event_client_constr.side_effect = Exception('All is well')
  93         lastHour = datetime.today() - timedelta(hours = 1)
  94         self._test_subscribe(
  95             poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
  96                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d %H:%M:%S')]])
  97
  98     def test_subscribe_websocket_with_start_time_next_hour(self):
  99         nextHour = datetime.today() + timedelta(hours = 1)
 100         with self.assertRaises(Queue.Empty):
 101             self._test_subscribe(
 102                 poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
 103                     additional_filters=[['created_at', '>=', nextHour.strftime('%Y-%m-%d %H:%M:%S')]], expected=0)
 104
 105     @mock.patch('arvados.events.EventClient.__init__')
 106     def test_subscribe_poll_with_start_time_next_hour(self, event_client_constr):
 107         event_client_constr.side_effect = Exception('All is well')
 108         nextHour = datetime.today() + timedelta(hours = 1)
 109         with self.assertRaises(Queue.Empty):
 110             self._test_subscribe(
 111                 poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
 112                     additional_filters=[['created_at', '>=', nextHour.strftime('%Y-%m-%d %H:%M:%S')]], expected=0)
 113
 114     def test_subscribe_websocket_with_start_time_tomorrow(self):
 115         tomorrow = datetime.today() + timedelta(hours = 24)
 116         with self.assertRaises(Queue.Empty):
 117             self._test_subscribe(
 118                 poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
 119                     additional_filters=[['created_at', '>=', tomorrow.strftime('%Y-%m-%d')]], expected=0)
 120
 121     @mock.patch('arvados.events.EventClient.__init__')
 122     def test_subscribe_poll_with_start_time_tomorrow(self, event_client_constr):
 123         event_client_constr.side_effect = Exception('All is well')
 124         tomorrow = datetime.today() + timedelta(hours = 24)
 125         with self.assertRaises(Queue.Empty):
 126             self._test_subscribe(
 127                 poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
 128                     additional_filters=[['created_at', '>=', tomorrow.strftime('%Y-%m-%d')]], expected=0)