Merge branch '4759-timestamp-precision-TC' closes #4759

[arvados.git] / sdk / python / tests / test_collections.py

diff --git a/sdk/python/tests/test_collections.py b/sdk/python/tests/test_collections.py
index 254a29f313eac8a898a18eea3a3f8b9348cd810f..dbbe3f5e73deca582b65f42900f8181e52a63a02 100644 (file)
--- a/sdk/python/tests/test_collections.py
+++ b/sdk/python/tests/test_collections.py
@@ -500,17 +500,7 @@ class ArvadosCollectionsTest(run_test_server.TestCaseWithServers,
                 ).manifest_text())
 
 
-class CollectionTestMixin(object):
-    PROXY_RESPONSE = {
-        'items_available': 1,
-        'items': [{
-                'uuid': 'zzzzz-bi6l4-mockproxy012345',
-                'owner_uuid': 'zzzzz-tpzed-mockowner012345',
-                'service_host': tutil.TEST_HOST,
-                'service_port': 65535,
-                'service_ssl_flag': True,
-                'service_type': 'proxy',
-                }]}
+class CollectionTestMixin(tutil.ApiClientMock):
     API_COLLECTIONS = run_test_server.fixture('collections')
     DEFAULT_COLLECTION = API_COLLECTIONS['foo_file']
     DEFAULT_DATA_HASH = DEFAULT_COLLECTION['portable_data_hash']
@@ -520,20 +510,9 @@ class CollectionTestMixin(object):
     ALT_DATA_HASH = ALT_COLLECTION['portable_data_hash']
     ALT_MANIFEST = ALT_COLLECTION['manifest_text']
 
-    def _mock_api_call(self, mock_method, code, body):
-        mock_method = mock_method().execute
-        if code == 200:
-            mock_method.return_value = body
-        else:
-            mock_method.side_effect = arvados.errors.ApiError(
-                tutil.fake_httplib2_response(code), "{}")
-
-    def mock_keep_services(self, api_mock, code, body):
-        self._mock_api_call(api_mock.keep_services().accessible, code, body)
-
-    def api_client_mock(self, code=200):
-        client = mock.MagicMock(name='api_client')
-        self.mock_keep_services(client, code, self.PROXY_RESPONSE)
+    def api_client_mock(self, status=200):
+        client = super(CollectionTestMixin, self).api_client_mock()
+        self.mock_keep_services(client, status=status, service_type='proxy', count=1)
         return client
 
 
@@ -543,9 +522,9 @@ class CollectionReaderTestCase(unittest.TestCase, CollectionTestMixin):
         body = self.API_COLLECTIONS.get(body)
         self._mock_api_call(api_mock.collections().get, code, body)
 
-    def api_client_mock(self, code=200):
-        client = super(CollectionReaderTestCase, self).api_client_mock(code)
-        self.mock_get_collection(client, code, 'foo_file')
+    def api_client_mock(self, status=200):
+        client = super(CollectionReaderTestCase, self).api_client_mock()
+        self.mock_get_collection(client, status, 'foo_file')
         return client
 
     def test_init_no_default_retries(self):
@@ -643,6 +622,20 @@ class CollectionReaderTestCase(unittest.TestCase, CollectionTestMixin):
                                           api_client=client)
         self.assertEqual('', reader.manifest_text())
 
+    def test_api_response(self):
+        client = self.api_client_mock()
+        reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+        self.assertEqual(self.DEFAULT_COLLECTION, reader.api_response())
+
+    def test_api_response_with_collection_from_keep(self):
+        client = self.api_client_mock()
+        self.mock_get_collection(client, 404, 'foo')
+        with tutil.mock_get_responses(self.DEFAULT_MANIFEST, 200):
+            reader = arvados.CollectionReader(self.DEFAULT_DATA_HASH,
+                                              api_client=client)
+            api_response = reader.api_response()
+        self.assertIsNone(api_response)
+
     def check_open_file(self, coll_file, stream_name, file_name, file_size):
         self.assertFalse(coll_file.closed, "returned file is not open")
         self.assertEqual(stream_name, coll_file.stream_name())
@@ -689,8 +682,8 @@ class CollectionWriterTestCase(unittest.TestCase, CollectionTestMixin):
         return tutil.mock_put_responses(body, *codes, **headers)
 
     def foo_writer(self, **kwargs):
-        api_client = self.api_client_mock()
-        writer = arvados.CollectionWriter(api_client, **kwargs)
+        kwargs.setdefault('api_client', self.api_client_mock())
+        writer = arvados.CollectionWriter(**kwargs)
         writer.start_new_file('foo')
         writer.write('foo')
         return writer
@@ -706,6 +699,32 @@ class CollectionWriterTestCase(unittest.TestCase, CollectionTestMixin):
             with self.assertRaises(arvados.errors.KeepWriteError):
                 writer.finish()
 
+    def test_write_insufficient_replicas_via_proxy(self):
+        writer = self.foo_writer(replication=3)
+        with self.mock_keep(None, 200, headers={'x-keep-replicas-stored': 2}):
+            with self.assertRaises(arvados.errors.KeepWriteError):
+                writer.manifest_text()
+
+    def test_write_insufficient_replicas_via_disks(self):
+        client = mock.MagicMock(name='api_client')
+        self.mock_keep_services(client, status=200, service_type='disk', count=2)
+        writer = self.foo_writer(api_client=client, replication=3)
+        with self.mock_keep(
+                None, 200, 200,
+                **{'x-keep-replicas-stored': 1}) as keepmock:
+            with self.assertRaises(arvados.errors.KeepWriteError):
+                writer.manifest_text()
+
+    def test_write_three_replicas(self):
+        client = mock.MagicMock(name='api_client')
+        self.mock_keep_services(client, status=200, service_type='disk', count=6)
+        writer = self.foo_writer(api_client=client, replication=3)
+        with self.mock_keep(
+                None, 500, 500, 500, 200, 200, 200,
+                **{'x-keep-replicas-stored': 1}) as keepmock:
+            writer.manifest_text()
+            self.assertEqual(6, keepmock.call_count)
+
     def test_write_whole_collection_through_retries(self):
         writer = self.foo_writer(num_retries=2)
         with self.mock_keep(self.DEFAULT_DATA_HASH,