+ def test_read_arbitrary_data_with_collection_reader(self):
+ # arv-get relies on this to do "arv-get {keep-locator} -".
+ self.write_foo_bar_baz()
+ self.assertEqual(
+ 'foobar',
+ arvados.CollectionReader(
+ '3858f62230ac3c915f300c664312c63f+6'
+ ).manifest_text())
+
+
+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']
+ DEFAULT_MANIFEST = DEFAULT_COLLECTION['manifest_text']
+ DEFAULT_UUID = DEFAULT_COLLECTION['uuid']
+ ALT_COLLECTION = API_COLLECTIONS['bar_file']
+ ALT_DATA_HASH = ALT_COLLECTION['portable_data_hash']
+ ALT_MANIFEST = ALT_COLLECTION['manifest_text']
+
+ 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
+
+
+@tutil.skip_sleep
+class CollectionReaderTestCase(unittest.TestCase, CollectionTestMixin):
+ def mock_get_collection(self, api_mock, code, body):
+ body = self.API_COLLECTIONS.get(body)
+ self._mock_api_call(api_mock.collections().get, code, body)
+
+ 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):
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+ reader.manifest_text()
+ client.collections().get().execute.assert_called_with(num_retries=0)
+
+ def test_uuid_init_success(self):
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client,
+ num_retries=3)
+ self.assertEqual(self.DEFAULT_COLLECTION['manifest_text'],
+ reader.manifest_text())
+ client.collections().get().execute.assert_called_with(num_retries=3)
+
+ def test_uuid_init_failure_raises_api_error(self):
+ client = self.api_client_mock(500)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+ with self.assertRaises(arvados.errors.ApiError):
+ reader.manifest_text()
+
+ def test_locator_init(self):
+ client = self.api_client_mock(200)
+ # Ensure Keep will not return anything if asked.
+ with tutil.mock_get_responses(None, 404):
+ reader = arvados.CollectionReader(self.DEFAULT_DATA_HASH,
+ api_client=client)
+ self.assertEqual(self.DEFAULT_MANIFEST, reader.manifest_text())
+
+ def test_locator_init_fallback_to_keep(self):
+ # crunch-job needs this to read manifests that have only ever
+ # been written to Keep.
+ client = self.api_client_mock(200)
+ self.mock_get_collection(client, 404, None)
+ with tutil.mock_get_responses(self.DEFAULT_MANIFEST, 200):
+ reader = arvados.CollectionReader(self.DEFAULT_DATA_HASH,
+ api_client=client)
+ self.assertEqual(self.DEFAULT_MANIFEST, reader.manifest_text())
+
+ def test_uuid_init_no_fallback_to_keep(self):
+ # Do not look up a collection UUID in Keep.
+ client = self.api_client_mock(404)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID,
+ api_client=client)
+ with tutil.mock_get_responses(self.DEFAULT_MANIFEST, 200):
+ with self.assertRaises(arvados.errors.ApiError):
+ reader.manifest_text()
+
+ def test_try_keep_first_if_permission_hint(self):
+ # To verify that CollectionReader tries Keep first here, we
+ # mock API server to return the wrong data.
+ client = self.api_client_mock(200)
+ with tutil.mock_get_responses(self.ALT_MANIFEST, 200):
+ self.assertEqual(
+ self.ALT_MANIFEST,
+ arvados.CollectionReader(
+ self.ALT_DATA_HASH + '+Affffffffffffffffffffffffffffffffffffffff@fedcba98',
+ api_client=client).manifest_text())
+
+ def test_init_num_retries_propagated(self):
+ # More of an integration test...
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client,
+ num_retries=3)
+ with tutil.mock_get_responses('foo', 500, 500, 200):
+ self.assertEqual('foo',
+ ''.join(f.read(9) for f in reader.all_files()))
+
+ def test_read_nonnormalized_manifest_with_collection_reader(self):
+ # client should be able to use CollectionReader on a manifest without normalizing it
+ client = self.api_client_mock(500)
+ nonnormal = ". acbd18db4cc2f85cedef654fccc4a4d8+3+Aabadbadbee@abeebdee 0:3:foo.txt 1:0:bar.txt 0:3:foo.txt\n"
+ reader = arvados.CollectionReader(
+ nonnormal,
+ api_client=client, num_retries=0)
+ # Ensure stripped_manifest() doesn't mangle our manifest in
+ # any way other than stripping hints.
+ self.assertEqual(
+ re.sub('\+[^\d\s\+]+', '', nonnormal),
+ reader.stripped_manifest())
+ # Ensure stripped_manifest() didn't mutate our reader.
+ self.assertEqual(nonnormal, reader.manifest_text())
+ # Ensure the files appear in the order given in the manifest.
+ self.assertEqual(
+ [[6, '.', 'foo.txt'],
+ [0, '.', 'bar.txt']],
+ [[f.size(), f.stream_name(), f.name()]
+ for f in reader.all_streams()[0].all_files()])
+
+ def test_read_empty_collection(self):
+ client = self.api_client_mock(200)
+ self.mock_get_collection(client, 200, 'empty')
+ reader = arvados.CollectionReader('d41d8cd98f00b204e9800998ecf8427e+0',
+ 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())
+ self.assertEqual(file_name, coll_file.name())
+ self.assertEqual(file_size, coll_file.size())
+
+ def test_open_collection_file_one_argument(self):
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+ cfile = reader.open('./foo')
+ self.check_open_file(cfile, '.', 'foo', 3)
+
+ def test_open_collection_file_two_arguments(self):
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+ cfile = reader.open('.', 'foo')
+ self.check_open_file(cfile, '.', 'foo', 3)
+
+ def test_open_deep_file(self):
+ coll_name = 'collection_with_files_in_subdir'
+ client = self.api_client_mock(200)
+ self.mock_get_collection(client, 200, coll_name)
+ reader = arvados.CollectionReader(
+ self.API_COLLECTIONS[coll_name]['uuid'], api_client=client)
+ cfile = reader.open('./subdir2/subdir3/file2_in_subdir3.txt')
+ self.check_open_file(cfile, './subdir2/subdir3', 'file2_in_subdir3.txt',
+ 32)
+
+ def test_open_nonexistent_stream(self):
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+ self.assertRaises(ValueError, reader.open, './nonexistent', 'foo')
+
+ def test_open_nonexistent_file(self):
+ client = self.api_client_mock(200)
+ reader = arvados.CollectionReader(self.DEFAULT_UUID, api_client=client)
+ self.assertRaises(ValueError, reader.open, '.', 'nonexistent')
+
+
+@tutil.skip_sleep
+class CollectionWriterTestCase(unittest.TestCase, CollectionTestMixin):
+ def mock_keep(self, body, *codes, **headers):
+ headers.setdefault('x-keep-replicas-stored', 2)
+ return tutil.mock_put_responses(body, *codes, **headers)
+
+ def foo_writer(self, **kwargs):
+ kwargs.setdefault('api_client', self.api_client_mock())
+ writer = arvados.CollectionWriter(**kwargs)
+ writer.start_new_file('foo')
+ writer.write('foo')
+ return writer
+
+ def test_write_whole_collection(self):
+ writer = self.foo_writer()
+ with self.mock_keep(self.DEFAULT_DATA_HASH, 200, 200):
+ self.assertEqual(self.DEFAULT_DATA_HASH, writer.finish())
+
+ def test_write_no_default(self):
+ writer = self.foo_writer()
+ with self.mock_keep(None, 500):
+ 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,
+ 500, 500, 200, 500, 500, 200):
+ self.assertEqual(self.DEFAULT_DATA_HASH, writer.finish())
+
+ def test_flush_data_retries(self):
+ writer = self.foo_writer(num_retries=2)
+ foo_hash = self.DEFAULT_MANIFEST.split()[1]
+ with self.mock_keep(foo_hash, 500, 200):
+ writer.flush_data()
+ self.assertEqual(self.DEFAULT_MANIFEST, writer.manifest_text())
+
+ def test_one_open(self):
+ client = self.api_client_mock()
+ writer = arvados.CollectionWriter(client)
+ with writer.open('out') as out_file:
+ self.assertEqual('.', writer.current_stream_name())
+ self.assertEqual('out', writer.current_file_name())
+ out_file.write('test data')
+ data_loc = hashlib.md5('test data').hexdigest() + '+9'
+ self.assertTrue(out_file.closed, "writer file not closed after context")
+ self.assertRaises(ValueError, out_file.write, 'extra text')
+ with self.mock_keep(data_loc, 200) as keep_mock:
+ self.assertEqual(". {} 0:9:out\n".format(data_loc),
+ writer.manifest_text())
+
+ def test_open_writelines(self):
+ client = self.api_client_mock()
+ writer = arvados.CollectionWriter(client)
+ with writer.open('six') as out_file:
+ out_file.writelines(['12', '34', '56'])
+ data_loc = hashlib.md5('123456').hexdigest() + '+6'
+ with self.mock_keep(data_loc, 200) as keep_mock:
+ self.assertEqual(". {} 0:6:six\n".format(data_loc),
+ writer.manifest_text())
+
+ def test_open_flush(self):
+ client = self.api_client_mock()
+ writer = arvados.CollectionWriter(client)
+ with writer.open('flush_test') as out_file:
+ out_file.write('flush1')
+ data_loc1 = hashlib.md5('flush1').hexdigest() + '+6'
+ with self.mock_keep(data_loc1, 200) as keep_mock:
+ out_file.flush()
+ out_file.write('flush2')
+ data_loc2 = hashlib.md5('flush2').hexdigest() + '+6'
+ with self.mock_keep(data_loc2, 200) as keep_mock:
+ self.assertEqual(". {} {} 0:12:flush_test\n".format(data_loc1,
+ data_loc2),
+ writer.manifest_text())
+
+ def test_two_opens_same_stream(self):
+ client = self.api_client_mock()
+ writer = arvados.CollectionWriter(client)
+ with writer.open('.', '1') as out_file:
+ out_file.write('1st')
+ with writer.open('.', '2') as out_file:
+ out_file.write('2nd')
+ data_loc = hashlib.md5('1st2nd').hexdigest() + '+6'
+ with self.mock_keep(data_loc, 200) as keep_mock:
+ self.assertEqual(". {} 0:3:1 3:3:2\n".format(data_loc),
+ writer.manifest_text())
+
+ def test_two_opens_two_streams(self):
+ client = self.api_client_mock()
+ writer = arvados.CollectionWriter(client)
+ with writer.open('file') as out_file:
+ out_file.write('file')
+ data_loc1 = hashlib.md5('file').hexdigest() + '+4'
+ with self.mock_keep(data_loc1, 200) as keep_mock:
+ with writer.open('./dir', 'indir') as out_file:
+ out_file.write('indir')
+ data_loc2 = hashlib.md5('indir').hexdigest() + '+5'
+ with self.mock_keep(data_loc2, 200) as keep_mock:
+ expected = ". {} 0:4:file\n./dir {} 0:5:indir\n".format(
+ data_loc1, data_loc2)
+ self.assertEqual(expected, writer.manifest_text())
+
+ def test_dup_open_fails(self):
+ client = self.api_client_mock()
+ writer = arvados.CollectionWriter(client)
+ file1 = writer.open('one')
+ self.assertRaises(arvados.errors.AssertionError, writer.open, 'two')
+