sdk/python/tests/test_websockets.py

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