X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/56b3e82324aea74cee083d3831c89eb7b553fb9c..cbf80c08daa5f9099d0821603a128967254709ed:/sdk/python/tests/test_collections.py

diff --git a/sdk/python/tests/test_collections.py b/sdk/python/tests/test_collections.py
index c991154e7f..4ca8dfef9e 100644
--- 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):
@@ -703,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
@@ -720,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, 200, 500, 200, 500, 200, 200,
+                **{'x-keep-replicas-stored': 1}) as keepmock:
+            writer.manifest_text()
+            self.assertEqual(5, 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,