+@tutil.skip_sleep
+class KeepClientCacheTestCase(unittest.TestCase, tutil.ApiClientMock):
+ def setUp(self):
+ self.api_client = self.mock_keep_services(count=2)
+ self.keep_client = arvados.KeepClient(api_client=self.api_client)
+ self.data = b'xyzzy'
+ self.locator = '1271ed5ef305aadabc605b1609e24c52'
+
+ @mock.patch('arvados.KeepClient.KeepService.get')
+ def test_get_request_cache(self, get_mock):
+ with tutil.mock_keep_responses(self.data, 200, 200):
+ self.keep_client.get(self.locator)
+ self.keep_client.get(self.locator)
+ # Request already cached, don't require more than one request
+ get_mock.assert_called_once()
+
+ @mock.patch('arvados.KeepClient.KeepService.get')
+ def test_head_request_cache(self, get_mock):
+ with tutil.mock_keep_responses(self.data, 200, 200):
+ self.keep_client.head(self.locator)
+ self.keep_client.head(self.locator)
+ # Don't cache HEAD requests so that they're not confused with GET reqs
+ self.assertEqual(2, get_mock.call_count)
+
+ @mock.patch('arvados.KeepClient.KeepService.get')
+ def test_head_and_then_get_return_different_responses(self, get_mock):
+ head_resp = None
+ get_resp = None
+ get_mock.side_effect = ['first response', 'second response']
+ with tutil.mock_keep_responses(self.data, 200, 200):
+ head_resp = self.keep_client.head(self.locator)
+ get_resp = self.keep_client.get(self.locator)
+ self.assertEqual('first response', head_resp)
+ # First reponse was not cached because it was from a HEAD request.
+ self.assertNotEqual(head_resp, get_resp)
+
+@tutil.skip_sleep
+class KeepStorageClassesTestCase(unittest.TestCase, tutil.ApiClientMock):
+ def setUp(self):
+ self.api_client = self.mock_keep_services(count=2)
+ self.keep_client = arvados.KeepClient(api_client=self.api_client)
+ self.data = b'xyzzy'
+ self.locator = '1271ed5ef305aadabc605b1609e24c52'
+
+ def test_multiple_default_storage_classes_req_header(self):
+ api_mock = self.api_client_mock()
+ api_mock.config.return_value = {
+ 'StorageClasses': {
+ 'foo': { 'Default': True },
+ 'bar': { 'Default': True },
+ 'baz': { 'Default': False }
+ }
+ }
+ api_client = self.mock_keep_services(api_mock=api_mock, count=2)
+ keep_client = arvados.KeepClient(api_client=api_client)
+ resp_hdr = {
+ 'x-keep-storage-classes-confirmed': 'foo=1, bar=1',
+ 'x-keep-replicas-stored': 1
+ }
+ with tutil.mock_keep_responses(self.locator, 200, **resp_hdr) as mock:
+ keep_client.put(self.data, copies=1)
+ req_hdr = mock.responses[0]
+ self.assertIn(
+ 'X-Keep-Storage-Classes: bar, foo', req_hdr.getopt(pycurl.HTTPHEADER))
+
+ def test_storage_classes_req_header(self):
+ self.assertEqual(
+ self.api_client.config()['StorageClasses'],
+ {'default': {'Default': True}})
+ cases = [
+ # requested, expected
+ [['foo'], 'X-Keep-Storage-Classes: foo'],
+ [['bar', 'foo'], 'X-Keep-Storage-Classes: bar, foo'],
+ [[], 'X-Keep-Storage-Classes: default'],
+ [None, 'X-Keep-Storage-Classes: default'],
+ ]
+ for req_classes, expected_header in cases:
+ headers = {'x-keep-replicas-stored': 1}
+ if req_classes is None or len(req_classes) == 0:
+ confirmed_hdr = 'default=1'
+ elif len(req_classes) > 0:
+ confirmed_hdr = ', '.join(["{}=1".format(cls) for cls in req_classes])
+ headers.update({'x-keep-storage-classes-confirmed': confirmed_hdr})
+ with tutil.mock_keep_responses(self.locator, 200, **headers) as mock:
+ self.keep_client.put(self.data, copies=1, classes=req_classes)
+ req_hdr = mock.responses[0]
+ self.assertIn(expected_header, req_hdr.getopt(pycurl.HTTPHEADER))
+
+ def test_partial_storage_classes_put(self):
+ headers = {
+ 'x-keep-replicas-stored': 1,
+ 'x-keep-storage-classes-confirmed': 'foo=1'}
+ with tutil.mock_keep_responses(self.locator, 200, 503, **headers) as mock:
+ with self.assertRaises(arvados.errors.KeepWriteError):
+ self.keep_client.put(self.data, copies=1, classes=['foo', 'bar'])
+ # 1st request, both classes pending
+ req1_headers = mock.responses[0].getopt(pycurl.HTTPHEADER)
+ self.assertIn('X-Keep-Storage-Classes: bar, foo', req1_headers)
+ # 2nd try, 'foo' class already satisfied
+ req2_headers = mock.responses[1].getopt(pycurl.HTTPHEADER)
+ self.assertIn('X-Keep-Storage-Classes: bar', req2_headers)
+
+ def test_successful_storage_classes_put_requests(self):
+ cases = [
+ # wanted_copies, wanted_classes, confirmed_copies, confirmed_classes, expected_requests
+ [ 1, ['foo'], 1, 'foo=1', 1],
+ [ 1, ['foo'], 2, 'foo=2', 1],
+ [ 2, ['foo'], 2, 'foo=2', 1],
+ [ 2, ['foo'], 1, 'foo=1', 2],
+ [ 1, ['foo', 'bar'], 1, 'foo=1, bar=1', 1],
+ [ 1, ['foo', 'bar'], 2, 'foo=2, bar=2', 1],
+ [ 2, ['foo', 'bar'], 2, 'foo=2, bar=2', 1],
+ [ 2, ['foo', 'bar'], 1, 'foo=1, bar=1', 2],
+ [ 1, ['foo', 'bar'], 1, None, 1],
+ [ 1, ['foo'], 1, None, 1],
+ [ 2, ['foo'], 2, None, 1],
+ [ 2, ['foo'], 1, None, 2],
+ ]
+ for w_copies, w_classes, c_copies, c_classes, e_reqs in cases:
+ headers = {'x-keep-replicas-stored': c_copies}
+ if c_classes is not None:
+ headers.update({'x-keep-storage-classes-confirmed': c_classes})
+ with tutil.mock_keep_responses(self.locator, 200, 200, **headers) as mock:
+ case_desc = 'wanted_copies={}, wanted_classes="{}", confirmed_copies={}, confirmed_classes="{}", expected_requests={}'.format(w_copies, ', '.join(w_classes), c_copies, c_classes, e_reqs)
+ self.assertEqual(self.locator,
+ self.keep_client.put(self.data, copies=w_copies, classes=w_classes),
+ case_desc)
+ self.assertEqual(e_reqs, mock.call_count, case_desc)
+
+ def test_failed_storage_classes_put_requests(self):
+ cases = [
+ # wanted_copies, wanted_classes, confirmed_copies, confirmed_classes, return_code
+ [ 1, ['foo'], 1, 'bar=1', 200],
+ [ 1, ['foo'], 1, None, 503],
+ [ 2, ['foo'], 1, 'bar=1, foo=0', 200],
+ [ 3, ['foo'], 1, 'bar=1, foo=1', 200],
+ [ 3, ['foo', 'bar'], 1, 'bar=2, foo=1', 200],
+ ]
+ for w_copies, w_classes, c_copies, c_classes, return_code in cases:
+ headers = {'x-keep-replicas-stored': c_copies}
+ if c_classes is not None:
+ headers.update({'x-keep-storage-classes-confirmed': c_classes})
+ with tutil.mock_keep_responses(self.locator, return_code, return_code, **headers):
+ case_desc = 'wanted_copies={}, wanted_classes="{}", confirmed_copies={}, confirmed_classes="{}"'.format(w_copies, ', '.join(w_classes), c_copies, c_classes)
+ with self.assertRaises(arvados.errors.KeepWriteError, msg=case_desc):
+ self.keep_client.put(self.data, copies=w_copies, classes=w_classes)
+
+@tutil.skip_sleep
+class KeepXRequestIdTestCase(unittest.TestCase, tutil.ApiClientMock):
+ def setUp(self):
+ self.api_client = self.mock_keep_services(count=2)
+ self.keep_client = arvados.KeepClient(api_client=self.api_client)
+ self.data = b'xyzzy'
+ self.locator = '1271ed5ef305aadabc605b1609e24c52'
+ self.test_id = arvados.util.new_request_id()
+ self.assertRegex(self.test_id, r'^req-[a-z0-9]{20}$')
+ # If we don't set request_id to None explicitly here, it will
+ # return <MagicMock name='api_client_mock.request_id'
+ # id='123456789'>:
+ self.api_client.request_id = None
+
+ def test_default_to_api_client_request_id(self):
+ self.api_client.request_id = self.test_id
+ with tutil.mock_keep_responses(self.locator, 200, 200) as mock:
+ self.keep_client.put(self.data)
+ self.assertEqual(2, len(mock.responses))
+ for resp in mock.responses:
+ self.assertProvidedRequestId(resp)
+
+ with tutil.mock_keep_responses(self.data, 200) as mock:
+ self.keep_client.get(self.locator)
+ self.assertProvidedRequestId(mock.responses[0])
+
+ with tutil.mock_keep_responses(b'', 200) as mock:
+ self.keep_client.head(self.locator)
+ self.assertProvidedRequestId(mock.responses[0])
+
+ def test_explicit_request_id(self):
+ with tutil.mock_keep_responses(self.locator, 200, 200) as mock:
+ self.keep_client.put(self.data, request_id=self.test_id)
+ self.assertEqual(2, len(mock.responses))
+ for resp in mock.responses:
+ self.assertProvidedRequestId(resp)
+
+ with tutil.mock_keep_responses(self.data, 200) as mock:
+ self.keep_client.get(self.locator, request_id=self.test_id)
+ self.assertProvidedRequestId(mock.responses[0])
+
+ with tutil.mock_keep_responses(b'', 200) as mock:
+ self.keep_client.head(self.locator, request_id=self.test_id)
+ self.assertProvidedRequestId(mock.responses[0])
+
+ def test_automatic_request_id(self):
+ with tutil.mock_keep_responses(self.locator, 200, 200) as mock:
+ self.keep_client.put(self.data)
+ self.assertEqual(2, len(mock.responses))
+ for resp in mock.responses:
+ self.assertAutomaticRequestId(resp)
+
+ with tutil.mock_keep_responses(self.data, 200) as mock:
+ self.keep_client.get(self.locator)
+ self.assertAutomaticRequestId(mock.responses[0])
+
+ with tutil.mock_keep_responses(b'', 200) as mock:
+ self.keep_client.head(self.locator)
+ self.assertAutomaticRequestId(mock.responses[0])
+
+ def assertAutomaticRequestId(self, resp):
+ hdr = [x for x in resp.getopt(pycurl.HTTPHEADER)
+ if x.startswith('X-Request-Id: ')][0]
+ self.assertNotEqual(hdr, 'X-Request-Id: '+self.test_id)
+ self.assertRegex(hdr, r'^X-Request-Id: req-[a-z0-9]{20}$')
+
+ def assertProvidedRequestId(self, resp):
+ self.assertIn('X-Request-Id: '+self.test_id,
+ resp.getopt(pycurl.HTTPHEADER))
+
+