def test_create_dataset(make_stubber, error_code):
lookoutvision_client = boto3.client('lookoutvision')
lookoutvision_stubber = make_stubber(lookoutvision_client)
project_name = 'test-project_name'
bucket = 'test-bucket'
object_key = 'test-object'
manifest_file = f'{bucket}/{object_key}'
dataset_type = 'train'
status = 'CREATE_COMPLETE'
message = 'Test message'
lookoutvision_stubber.stub_create_dataset(project_name,
dataset_type,
bucket,
object_key,
status,
message,
error_code=error_code)
if error_code is None:
lookoutvision_stubber.stub_describe_dataset(project_name, dataset_type,
status, message)
if error_code is None:
Datasets.create_dataset(lookoutvision_client, project_name,
manifest_file, dataset_type)
else:
with pytest.raises(ClientError) as exc_info:
Datasets.create_dataset(lookoutvision_client, project_name,
manifest_file, dataset_type)
assert exc_info.value.response['Error']['Code'] == error_code
def create_dataset(lookoutvision_client, s3_resource, bucket, project_name,
dataset_images, dataset_type):
"""
Creates a manifest from images in the supplied bucket and then creates
a dataset.
:param lookoutvision_client: A Boto3 Lookout for Vision client.
:param s3_resource: A Boto3 Amazon S3 client.
:param bucket: The bucket that stores the manifest file.
:param project_name: The project in which to create the dataset.
:param dataset_images: The location of the images referenced by the dataset.
:param dataset_type: The type of dataset to create (train or test).
"""
print(f"Creating {dataset_type} dataset...")
manifest_file = f"s3://{bucket}/{project_name}/manifests/{dataset_type}.manifest"
logger.info("Creating %s manifest file in %s.", dataset_type,
manifest_file)
Datasets.create_manifest_file_s3(s3_resource, dataset_images,
manifest_file)
logger.info("Create %s dataset for project %s", dataset_type, project_name)
Datasets.create_dataset(lookoutvision_client, project_name, manifest_file,
dataset_type)
