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_websocket_with_start_time_date_only(self, event_client_constr):
  78         lastHour = datetime.today() - timedelta(hours = 1)
  79         self._test_subscribe(
  80             poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
  81                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d')]])
  82
  83     def test_subscribe_websocket_with_start_time_last_hour(self):
  84         lastHour = datetime.today() - timedelta(hours = 1)
  85         self._test_subscribe(
  86             poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
  87                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d %H:%M:%S')]])
  88
  89     @mock.patch('arvados.events.EventClient.__init__')
  90     def test_subscribe_poll_with_start_time_last_hour(self, event_client_constr):
  91         lastHour = datetime.today() - timedelta(hours = 1)
  92         self._test_subscribe(
  93             poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
  94                 additional_filters=[['created_at', '>=', lastHour.strftime('%Y-%m-%d %H:%M:%S')]])
  95
  96     def test_subscribe_websocket_with_start_time_next_hour(self):
  97         nextHour = datetime.today() + timedelta(hours = 1)
  98         with self.assertRaises(Queue.Empty):
  99             self._test_subscribe(
 100                 poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
 101                     additional_filters=[['created_at', '>=', nextHour.strftime('%Y-%m-%d %H:%M:%S')]], expected=0)
 102
 103     @mock.patch('arvados.events.EventClient.__init__')
 104     def test_subscribe_poll_with_start_time_next_hour(self, event_client_constr):
 105         nextHour = datetime.today() + timedelta(hours = 1)
 106         with self.assertRaises(Queue.Empty):
 107             self._test_subscribe(
 108                 poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
 109                     additional_filters=[['created_at', '>=', nextHour.strftime('%Y-%m-%d %H:%M:%S')]], expected=0)
 110
 111     def test_subscribe_websocket_with_start_time_tomorrow(self):
 112         tomorrow = datetime.today() + timedelta(hours = 24)
 113         with self.assertRaises(Queue.Empty):
 114             self._test_subscribe(
 115                 poll_fallback=False, expect_type=arvados.events.EventClient, last_log_id=1,
 116                     additional_filters=[['created_at', '>=', tomorrow.strftime('%Y-%m-%d')]], expected=0)
 117
 118     @mock.patch('arvados.events.EventClient.__init__')
 119     def test_subscribe_poll_with_start_time_tomorrow(self, event_client_constr):
 120         tomorrow = datetime.today() + timedelta(hours = 24)
 121         with self.assertRaises(Queue.Empty):
 122             self._test_subscribe(
 123                 poll_fallback=1, expect_type=arvados.events.PollClient, last_log_id=1,
 124                     additional_filters=[['created_at', '>=', tomorrow.strftime('%Y-%m-%d')]], expected=0)