Merge branch 'master' into 6473-fetch-events-starting-at
[arvados.git] / 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)