Fix prefixes and source URIs (#57)

* Fix prefixes and source URIs

* Update test

Author: di

main.py

@@ -124,7 +124,7 @@ def load_processed_files_into_bigquery(event, context):
         # Otherwise, this was triggered via cron, use the current time
         partition = datetime.datetime.utcnow().strftime("%Y%m%d")
 
-    folder = f"gs://{RESULT_BUCKET}/processed/{partition}"
+    folder = f"processed/{partition}"
 
     # Load the data into the dataset(s)
     job_config = bigquery.LoadJobConfig()
@@ -137,10 +137,16 @@ def load_processed_files_into_bigquery(event, context):
     bigquery_client = bigquery.Client()
 
     # Get the processed files we're loading
-    download_prefix = f"{folder}/downloads-*.json"
-    download_source_uris = bucket.list_blobs(prefix=download_prefix)
-    simple_prefix = f"{folder}/simple-*.json"
-    simple_source_uris = bucket.list_blobs(prefix=simple_prefix)
+    download_prefix = f"{folder}/downloads-"
+    download_source_blobs = bucket.list_blobs(prefix=download_prefix)
+    download_source_uris = [
+        f"gs://{blob.bucket.name}/{blob.name}" for blob in download_source_blobs
+    ]
+    simple_prefix = f"{folder}/simple-"
+    simple_source_blobs = bucket.list_blobs(prefix=simple_prefix)
+    simple_source_uris = [
+        f"gs://{blob.bucket.name}/{blob.name}" for blob in simple_source_blobs
+    ]
 
     for DATASET in DATASETS:
         dataset_ref = bigquery.dataset.DatasetReference.from_string(

test_functions.py

@@ -147,8 +147,12 @@ def test_load_processed_files_into_bigquery(
 
     blob_lists = {}
 
+    bucket = pretend.stub(name=RESULT_BUCKET)
+
+    blob_stub = pretend.stub(name="blobname", bucket=bucket)
+
     def _generate_blob_list(prefix):
-        blob_list = pretend.stub(prefix=prefix)
+        blob_list = [blob_stub]
         blob_lists[prefix] = blob_list
         return blob_list
 
@@ -199,17 +203,14 @@ def _generate_blob_list(prefix):
         pretend.call(RESULT_BUCKET),
     ]
     assert bucket_stub.list_blobs.calls == [
-        pretend.call(
-            prefix=f"gs://my-result-bucket/processed/{partition}/downloads-*.json"
-        ),
-        pretend.call(
-            prefix=f"gs://my-result-bucket/processed/{partition}/simple-*.json"
-        ),
+        pretend.call(prefix=f"processed/{partition}/downloads-"),
+        pretend.call(prefix=f"processed/{partition}/simple-"),
     ]
     assert (
         load_job_stub.result.calls
         == [pretend.call()] * len(bigquery_dataset.split()) * 2
     )
-    assert set(bucket_stub.delete_blobs.calls) == set(
-        pretend.call(blobs=blobs) for blobs in blob_lists.values()
+    assert (
+        bucket_stub.delete_blobs.calls
+        == [pretend.call(blobs=[f"gs://{RESULT_BUCKET}/{blob_stub.name}"])] * 2
     )