def test_record_status_in_datastore_success(self):
num_files = 5
fnames = self._get_file_names(num_files)
entities = list()
for i in range(num_files):
key = KeyStub(ds.DATASTORE_PHOTO, fnames[i])
entities.append(GCDEntity(key=key))
client = Mock()
client.get_multi = Mock(return_value=entities)
uploader.datastore.Client = Mock(return_value=client)
uploader.datastore.key.Key = KeyStub
for success in (True, False):
# Call under test
ret_val = uploader._record_status_in_datastore(fnames,
success=success)
keys = [
KeyStub(ds.DATASTORE_PHOTO, f, project=config.PROJECT_ID)
for f in fnames
]
self.assertEqual(ret_val, [])
client.get_multi.assert_called_with(keys)
client.put_multi.assert_called_with(entities)
# Make sure the entities were updated correctly
for e in entities:
field = 'in_gcs' if success else 'gcs_upload_failed'
self.assertTrue(e[field])
def test_record_status_in_datastore_missing_entities_success(self):
num_files = 5
fnames = self._get_file_names(num_files)
entities = list()
for i in range(num_files / 2):
key = KeyStub(ds.DATASTORE_PHOTO, fnames[i])
entities.append(GCDEntity(key=key))
client = Mock()
client.get_multi = Mock(return_value=entities)
uploader.datastore.entity.Entity = GCDEntity
uploader.datastore.Client = Mock(return_value=client)
uploader.datastore.key.Key = KeyStub
# Call under test
ret_val = uploader._record_status_in_datastore(fnames, success=True)
keys = [
KeyStub(ds.DATASTORE_PHOTO, f, project=config.PROJECT_ID)
for f in fnames
]
self.assertEqual(ret_val, [])
client.get_multi.assert_called_with(keys)
client.put_multi.assert_called_with(entities)
# Make sure the entities were updated correctly
self.assertEqual(len(entities), len(fnames))
for e in entities:
self.assertTrue(e['in_gcs'])
def test_record_status_in_datastore_error_saving_to_datastore(self):
num_files = 5
fnames = self._get_file_names(num_files)
gcloud_error = GCloudError('')
# Must set the code since __str__ will be called on this instance.
# There is no way to set the code through the __init__ method
gcloud_error.code = 500
entities = [dict() for _ in range(num_files)]
client = Mock()
client.get_multi = Mock(return_value=entities)
client.put_multi = Mock(side_effect=gcloud_error)
uploader.datastore.Client = Mock(return_value=client)
uploader.datastore.key.Key = KeyStub
# Call under test
ret_val = uploader._record_status_in_datastore(fnames, success=True)
keys = [
KeyStub(ds.DATASTORE_PHOTO, f, project=config.PROJECT_ID)
for f in fnames
]
self.assertEqual(ret_val, fnames)
client.get_multi.assert_called_with(keys)
client.put_multi.assert_called_with(entities)
# Make sure the entities were updated correctly
for e in entities:
self.assertTrue(e['in_gcs'])
def test_record_status_in_datastore_restricted_fields_in_fetched_entities(
self):
num_files = 5
fnames = self._get_file_names(num_files)
# user is a restricted field for datastore photo entities - see
# common/constants.py
entities = [{'user': 123456789} for _ in range(num_files)]
client = Mock()
client.get_multi = Mock(return_value=entities)
uploader.datastore.Client = Mock(return_value=client)
uploader.datastore.key.Key = KeyStub
# Call under test
ret_val = uploader._record_status_in_datastore(fnames, success=True)
keys = [
KeyStub(ds.DATASTORE_PHOTO, f, project=config.PROJECT_ID)
for f in fnames
]
self.assertEqual(ret_val, list())
client.get_multi.assert_called_with(keys)
# Entities should have been put into datastore, despite having
# restricted fields - this is because the validate data function
# is only called on the new data.
client.put_multi.assert_called_with(entities)
def test_insert_missing_entities(self):
now = 'now'
fnames = self._get_file_names(10)
entities = list()
# Create entities for half the files
for i in range(len(fnames)):
if i % 2 != 1:
key = KeyStub(ds.DATASTORE_PHOTO, fnames[i])
entities.append(GCDEntity(key=key))
uploader.datetime.now = Mock(return_value=now)
uploader.datastore.entity.Entity = GCDEntity
uploader.datastore.key.Key = KeyStub
# Insert missing entities should work with file paths or names
prefixes = ('', '/some/arbitrary/path/')
for prefix in prefixes:
fpaths = [prefix + fname for fname in fnames]
# Call under test
ret_val = uploader._insert_missing_entities(entities, fpaths)
self.assertEqual(len(ret_val), len(fnames))
for fname in fnames:
self.assertTrue(
util.in_list(ret_val, fname, key=lambda x: x.key.name))
# All the new eitities (the second half of them) should have the
# following
for i in range(len(fnames) / 2, len(fnames)):
self.assertNotIn('user', entities[i])
self.assertEqual(entities[i]['uploaded_date'], now)
def test_record_status_in_datastore_error_getting_entities(self):
fnames = self._get_file_names(5)
gcloud_error = GCloudError('')
# Must set the code since __str__ will be called on this instance.
# There is no way to set the code through the __init__ method
gcloud_error.code = 500
client = Mock()
client.get_multi = Mock(side_effect=gcloud_error)
uploader.datastore.Client = Mock(return_value=client)
uploader.datastore.key.Key = KeyStub
# Call under test
ret_val = uploader._record_status_in_datastore(fnames, success=True)
keys = [
KeyStub(ds.DATASTORE_PHOTO, f, project=config.PROJECT_ID)
for f in fnames
]
self.assertEqual(ret_val, fnames)
client.get_multi.assert_called_with(keys)
client.put_multi.assert_not_called()