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()