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_record_status_in_datastore_could_not_obtain_credentials(self): fnames = self._get_file_names(5) uploader.datastore.Client = Mock( side_effect=CouldNotObtainCredentialsError) # Call under test ret_val = uploader._record_status_in_datastore(fnames, success=True) self.assertEqual(ret_val, fnames)
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()