Updated tests to handle json lines format

Author: robbailiff2

scripts/manual_uploads/manual_s3_dynamo_upload.py

@@ -4,12 +4,11 @@
 import os
 import argparse
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union, Generator
 from decimal import Decimal
 
 
-
-def map_dynamo_type(value: Any):
+def map_dynamo_type(value: Any) -> Dict[str, Any]:
     if isinstance(value, str):
         return {"S": value}
     elif isinstance(value, bool):
@@ -29,7 +28,20 @@ def map_dynamo_type(value: Any):
         return {"S": value}
 
 
-def upload_to_s3(s3_client, bucket, filepath, dry_run=False):
+def load_json_lines(filepath: Union[str, Path]) -> Generator[Dict[str, Any], None, None]:
+    with open(filepath) as f:
+        for line in f:
+            if line.strip():
+                yield json.loads(line)
+
+
+def upload_to_s3(
+    s3_client: Any,
+    bucket: str,
+    filepath: Union[str, Path],
+    dry_run: bool = False
+) -> None:
+
     filename = os.path.basename(filepath)
     print(f"Filepath: {filepath}")
     s3_key = f"manual-uploads/{filename}"
@@ -45,21 +57,23 @@ def upload_to_s3(s3_client, bucket, filepath, dry_run=False):
         print(f"Failed to upload {filepath}: {e}")
 
 
-def upload_to_dynamo(dynamo_client, table_name, filepath):
-    with open(filepath) as f:
-        item = json.load(f)
+def upload_to_dynamo(
+    dynamo_client: Any,
+    table_name: str,
+    filepath: Union[str, Path],
+) -> None:
 
-    try:
-        dynamo_client.put_item(
-            TableName=table_name, Item={key: map_dynamo_type(value) for key, value in item.items()}
-        )
-        print(f"Uploaded {filepath} to DynamoDB table {table_name}")
-    except Exception as e:
-        print(f"Failed to upload {filepath}: {e}")
+    for item in load_json_lines(filepath):
+        try:
+            dynamo_client.put_item(
+                TableName=table_name, Item={key: map_dynamo_type(value) for key, value in item.items()}
+            )
+            print(f"Uploaded {filepath} to DynamoDB table {table_name}")
+        except Exception as e:
+            print(f"Failed to upload {filepath}: {e}")
 
 
-def run_upload(args=None):
-    print("\n\n\n***** We are in main *****\n\n\n")
+def run_upload(args: Optional[List[str]] = None) -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--env")
     parser.add_argument("--upload-s3", type=Path)

tests/utils/manual_s3_dynamo_upload_tests.py

@@ -8,20 +8,26 @@
 import boto3
 from moto import mock_aws
 
-from scripts.manual_uploads.manual_s3_dynamo_upload import run_upload
+from scripts.manual_uploads.manual_s3_dynamo_upload import run_upload, map_dynamo_type
 
 
 @pytest.fixture
 def test_data_dir(tmp_path):
-    data = {"ID_NUMBER": "123", "ATTRIBUTE_TYPE": "Test", "value": 99}
+    # {"NHS_NUMBER": "C456", "ATTRIBUTE_TYPE": "COHORTS", "COHORT_MEMBERSHIPS": [{"COHORT_LABEL": "under_75", "DATE_JOINED": "2025-01-01"},{"COHORT_LABEL": "over_75", "DATE_JOINED": "2025-01-01"}], "ACTION_FLAG": "ADD", "HASH": "fake_hash_delete_sim", "sync_time": "2025-05-01 12:00:00"}
+    data = [
+        {"NHS_NUMBER": "1234567890", "ATTRIBUTE_TYPE": "COHORTS", "COHORT_MEMBERSHIPS": [{"COHORT_LABEL": "under_75", "DATE_JOINED": "2025-01-01"}]},
+        {"NHS_NUMBER": "2345678901", "ATTRIBUTE_TYPE": "COHORTS", "COHORT_MEMBERSHIPS": [{"COHORT_LABEL": "over_75", "DATE_JOINED": "2025-01-01"}]},
+        {"NHS_NUMBER": "3456789012", "ATTRIBUTE_TYPE": "COHORTS", "COHORT_MEMBERSHIPS": [{"COHORT_LABEL": "16+_covid", "DATE_JOINED": "2025-01-01"}]}
+    ]
     file_path = tmp_path / "test.json"
     with open(file_path, "w") as f:
-        json.dump(data, f)
+        for item in data:
+            f.write(json.dumps(item) + "\n")
     return tmp_path, data
 
 
 @mock_aws
-def test_script_cli_end_to_end(test_data_dir, capsys):
+def test_script_cli_end_to_end(test_data_dir):
     # Arrange
     data_dir, expected_data = test_data_dir
     env = "test"
@@ -39,41 +45,44 @@ def test_script_cli_end_to_end(test_data_dir, capsys):
     dynamodb.create_table(
         TableName=dynamo_table,
         KeySchema=[
-            {"AttributeName": "ID_NUMBER", "KeyType": "HASH"},
+            {"AttributeName": "NHS_NUMBER", "KeyType": "HASH"},
             {"AttributeName": "ATTRIBUTE_TYPE", "KeyType": "RANGE"}
         ],
         AttributeDefinitions=[
-            {"AttributeName": "ID_NUMBER", "AttributeType": "S"},
+            {"AttributeName": "NHS_NUMBER", "AttributeType": "S"},
             {"AttributeName": "ATTRIBUTE_TYPE", "AttributeType": "S"}
         ],
         BillingMode="PAY_PER_REQUEST"
     )
 
     # Act
-    return_code = run_upload([
+    run_upload([
         "--env", env,
         "--upload-s3", str(data_dir),
         "--upload-dynamo", str(data_dir),
         "--region", region,
         "--s3-bucket", s3_bucket,
         "--dynamo-table", dynamo_table
     ])
-    captured = capsys.readouterr()
 
-    # Assert
     key = f"manual-uploads/test.json"
     obj = s3.get_object(Bucket=s3_bucket, Key=key)
-    uploaded_s3_data = json.load(obj["Body"])
-    assert uploaded_s3_data == expected_data
+    body = obj["Body"].read().decode("utf-8")
+    uploaded_s3_data = [json.loads(line) for line in body.splitlines() if line.strip()]
 
-    item = dynamodb.get_item(
-        TableName=dynamo_table,
-        Key={
-            "ID_NUMBER": {"S": expected_data["ID_NUMBER"]},
-            "ATTRIBUTE_TYPE": {"S": expected_data["ATTRIBUTE_TYPE"]}
+    dynamo_items = []
+    for expected_item in expected_data:
+        key = {
+            "NHS_NUMBER": {"S": expected_item["NHS_NUMBER"]},
+            "ATTRIBUTE_TYPE": {"S": expected_item["ATTRIBUTE_TYPE"]}
         }
-    )["Item"]
-    assert item["value"]["N"] == "99"
+        response = dynamodb.get_item(TableName=dynamo_table, Key=key)
+        item = response.get("Item")
+        assert item is not None, f"Missing item for key {key}"
+        dynamo_items.append(item)
 
-    assert "Uploaded" in captured.out
-    assert "Error" not in captured.err
+    expected_dynamo_items = [{k: map_dynamo_type(v) for k, v in item.items()} for item in expected_data]
+
+    # Assert
+    assert uploaded_s3_data == expected_data
+    assert dynamo_items == expected_dynamo_items