Merge branch '14259-pysdk-remote-block-copy'
[arvados.git] / sdk / python / tests / test_keep_client.py
index 35070f3f1c614e038eb27ec20c18731fd9f5c4cb..192a21cd66b4af3cc062b72345263b5afb7ccdef 100644 (file)
@@ -1,3 +1,7 @@
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
 from __future__ import absolute_import
 from __future__ import division
 from future import standard_library
@@ -13,13 +17,13 @@ import random
 import re
 import socket
 import sys
-import threading
 import time
 import unittest
 import urllib.parse
 
 import arvados
 import arvados.retry
+import arvados.util
 from . import arvados_testutil as tutil
 from . import keepstub
 from . import run_test_server
@@ -315,6 +319,48 @@ class KeepClientServiceTestCase(unittest.TestCase, tutil.ApiClientMock):
         self.assertEqual('100::1', service.hostname)
         self.assertEqual(10, service.port)
 
+    def test_insecure_disables_tls_verify(self):
+        api_client = self.mock_keep_services(count=1)
+        force_timeout = socket.timeout("timed out")
+
+        api_client.insecure = True
+        with tutil.mock_keep_responses(b'foo', 200) as mock:
+            keep_client = arvados.KeepClient(api_client=api_client)
+            keep_client.get('acbd18db4cc2f85cedef654fccc4a4d8+3')
+            self.assertEqual(
+                mock.responses[0].getopt(pycurl.SSL_VERIFYPEER),
+                0)
+
+        api_client.insecure = False
+        with tutil.mock_keep_responses(b'foo', 200) as mock:
+            keep_client = arvados.KeepClient(api_client=api_client)
+            keep_client.get('acbd18db4cc2f85cedef654fccc4a4d8+3')
+            # getopt()==None here means we didn't change the
+            # default. If we were using real pycurl instead of a mock,
+            # it would return the default value 1.
+            self.assertEqual(
+                mock.responses[0].getopt(pycurl.SSL_VERIFYPEER),
+                None)
+
+    def test_refresh_signature(self):
+        blk_digest = '6f5902ac237024bdd0c176cb93063dc4+11'
+        blk_sig = 'da39a3ee5e6b4b0d3255bfef95601890afd80709@53bed294'
+        local_loc = blk_digest+'+A'+blk_sig
+        remote_loc = blk_digest+'+R'+blk_sig
+        api_client = self.mock_keep_services(count=1)
+        headers = {'X-Keep-Locator':local_loc}
+        with tutil.mock_keep_responses('', 200, **headers):
+            # Check that the translated locator gets returned
+            keep_client = arvados.KeepClient(api_client=api_client)
+            self.assertEqual(local_loc, keep_client.refresh_signature(remote_loc))
+            # Check that refresh_signature() uses the correct method and headers
+            keep_client._get_or_head = mock.MagicMock()
+            keep_client.refresh_signature(remote_loc)
+            args, kwargs = keep_client._get_or_head.call_args_list[0]
+            self.assertIn(remote_loc, args)
+            self.assertEqual("HEAD", kwargs['method'])
+            self.assertIn('X-Keep-Signature', kwargs['headers'])
+
     # test_*_timeout verify that KeepClient instructs pycurl to use
     # the appropriate connection and read timeouts. They don't care
     # whether pycurl actually exhibits the expected timeout behavior
@@ -514,6 +560,77 @@ class KeepClientServiceTestCase(unittest.TestCase, tutil.ApiClientMock):
         self.assertEqual(1, req_mock.call_count)
 
 
+@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))
+
+
 @tutil.skip_sleep
 class KeepClientRendezvousTestCase(unittest.TestCase, tutil.ApiClientMock):
 
@@ -669,7 +786,7 @@ class KeepClientRendezvousTestCase(unittest.TestCase, tutil.ApiClientMock):
         self.check_64_zeros_error_order('put', arvados.errors.KeepWriteError)
 
 
-class KeepClientTimeout(unittest.TestCase, tutil.ApiClientMock):
+class KeepClientTimeout(keepstub.StubKeepServers, unittest.TestCase):
     # BANDWIDTH_LOW_LIM must be less than len(DATA) so we can transfer
     # 1s worth of data and then trigger bandwidth errors before running
     # out of data.
@@ -703,24 +820,6 @@ class KeepClientTimeout(unittest.TestCase, tutil.ApiClientMock):
             delta = round(time.time() - self.t0, 3)
             self.assertGreaterEqual(delta, self.tmin)
 
-    def setUp(self):
-        sock = socket.socket()
-        sock.bind(('0.0.0.0', 0))
-        self.port = sock.getsockname()[1]
-        sock.close()
-        self.server = keepstub.Server(('0.0.0.0', self.port), keepstub.Handler)
-        self.thread = threading.Thread(target=self.server.serve_forever)
-        self.thread.daemon = True # Exit thread if main proc exits
-        self.thread.start()
-        self.api_client = self.mock_keep_services(
-            count=1,
-            service_host='localhost',
-            service_port=self.port,
-        )
-
-    def tearDown(self):
-        self.server.shutdown()
-
     def keepClient(self, timeouts=(0.1, TIMEOUT_TIME, BANDWIDTH_LOW_LIM)):
         return arvados.KeepClient(
             api_client=self.api_client,
@@ -1114,7 +1213,7 @@ class AvoidOverreplication(unittest.TestCase, tutil.ApiClientMock):
 
         def finished(self):
             return False
-    
+
     def setUp(self):
         self.copies = 3
         self.pool = arvados.KeepClient.KeepWriterThreadPool(
@@ -1158,7 +1257,7 @@ class AvoidOverreplication(unittest.TestCase, tutil.ApiClientMock):
             self.pool.add_task(ks, None)
         self.pool.join()
         self.assertEqual(self.pool.done(), self.copies-1)
-    
+
 
 @tutil.skip_sleep
 class RetryNeedsMultipleServices(unittest.TestCase, tutil.ApiClientMock):
@@ -1193,3 +1292,29 @@ class RetryNeedsMultipleServices(unittest.TestCase, tutil.ApiClientMock):
             with self.assertRaises(arvados.errors.KeepWriteError):
                 self.keep_client.put('foo', num_retries=1, copies=2)
         self.assertEqual(2, req_mock.call_count)
+
+class KeepClientAPIErrorTest(unittest.TestCase):
+    def test_api_fail(self):
+        class ApiMock(object):
+            def __getattr__(self, r):
+                if r == "api_token":
+                    return "abc"
+                elif r == "insecure":
+                    return False
+                else:
+                    raise arvados.errors.KeepReadError()
+        keep_client = arvados.KeepClient(api_client=ApiMock(),
+                                             proxy='', local_store='')
+
+        # The bug this is testing for is that if an API (not
+        # keepstore) exception is thrown as part of a get(), the next
+        # attempt to get that same block will result in a deadlock.
+        # This is why there are two get()s in a row.  Unfortunately,
+        # the failure mode for this test is that the test suite
+        # deadlocks, there isn't a good way to avoid that without
+        # adding a special case that has no use except for this test.
+
+        with self.assertRaises(arvados.errors.KeepReadError):
+            keep_client.get("acbd18db4cc2f85cedef654fccc4a4d8+3")
+        with self.assertRaises(arvados.errors.KeepReadError):
+            keep_client.get("acbd18db4cc2f85cedef654fccc4a4d8+3")