Add CLI options to set test data AIP sizing (#352)

artefactual-labs · Nov 19, 2024 · a367239 · a367239
1 parent 911875f
commit a367239
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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