diff --git a/tools/generate-test-data b/tools/generate-test-data index 79524f09..fe723829 100755 --- a/tools/generate-test-data +++ b/tools/generate-test-data @@ -61,6 +61,20 @@ from config import CONFIGS help="Maximum number of possible format types to assign to files (default 50).", type=int, ) +@click.option( + "--minimum-aip-file-size", + "-f", + default=5000000, + help="Minimum size of AIP (in bytes: default 5000000).", + type=int, +) +@click.option( + "--maximum-aip-file-size", + "-g", + default=5000000000, + help="Maximum size of AIP (in bytes: default 5000000000).", + type=int, +) @click.option("--seed", default=0) def main( storage_services_to_create, @@ -70,6 +84,8 @@ def main( aip_min_file_count, aip_max_file_count, number_of_format_types, + minimum_aip_file_size, + maximum_aip_file_size, seed, ): # Initialize Flash app context @@ -145,8 +161,14 @@ def main( for _ in range(0, aips_to_create): aip = data.create_fake_aip( - pipeline.id, ss_id, sl.id, fetch_jobs[ss.id] + pipeline.id, + ss_id, + sl.id, + fetch_jobs[ss.id], + minimum_aip_file_size, + maximum_aip_file_size, ) + data.create_fake_aip_files( aip_min_file_count, aip_max_file_count, diff --git a/tools/helpers/data.py b/tools/helpers/data.py index a33a42de..2f8a7a2f 100644 --- a/tools/helpers/data.py +++ b/tools/helpers/data.py @@ -99,13 +99,20 @@ def create_fake_location(storage_service_id): return location -def create_fake_aip(pipeline_id, storage_service_id, storage_location_id, fetch_job_id): +def create_fake_aip( + pipeline_id, + storage_service_id, + storage_location_id, + fetch_job_id, + min_size, + max_size, +): aip = AIP( uuid=fake.uuid4(), transfer_name=fake.text(20)[:-1], create_date=fake.date_time_between(start_date="-5y"), mets_sha256=fake.sha256(), - size=randint(10000, 100_000_000), + size=randint(min_size, max_size), storage_service_id=storage_service_id, storage_location_id=storage_location_id, fetch_job_id=fetch_job_id, diff --git a/tools/tests/test_data.py b/tools/tests/test_data.py index be823fdf..77ca39a4 100644 --- a/tools/tests/test_data.py +++ b/tools/tests/test_data.py @@ -64,7 +64,7 @@ def test_create_fake_location(mock_db_add): def test_create_fake_aip(mock_db_add): - aip = data.create_fake_aip(1, 2, 3, 4) + aip = data.create_fake_aip(1, 2, 3, 4, 100, 100) assert aip.uuid assert type(aip.uuid) is str @@ -86,3 +86,4 @@ def test_create_fake_aip(mock_db_add): assert aip.storage_location_id == 3 assert aip.fetch_job_id == 4 assert aip.origin_pipeline_id == 1 + assert aip.size == 100