Merge pull request #146 from nationalarchives/fix/multi-message

Fix/multi message

ds-caselaw-ingester/lambda_function.py

@@ -53,6 +53,7 @@ class Message(object):
     def from_event(cls, event):
         decoder = json.decoder.JSONDecoder()
         message = decoder.decode(event["Records"][0]["Sns"]["Message"])
+        # passes a messagedict to the class
         return cls.from_message(message)
 
     @classmethod
@@ -155,6 +156,15 @@ class DocumentInsertionError(ReportableException):
     pass
 
 
+def all_messages(event) -> List[Message]:
+    """All the messages in the SNS event, as Message subclasses"""
+    decoder = json.decoder.JSONDecoder()
+    messages_as_decoded_json = [
+        decoder.decode(record["Sns"]["Message"]) for record in event["Records"]
+    ]
+    return [Message.from_message(message) for message in messages_as_decoded_json]
+
+
 def extract_xml_file(tar: tarfile.TarFile, xml_file_name: str):
     xml_file = None
     if xml_file_name:
@@ -429,9 +439,8 @@ def unpublish_updated_judgment(uri):
     api_client.set_published(uri, False)
 
 
-@rollbar.lambda_function
-def handler(event, context):
-    message = Message.from_event(event)
+def process_message(message):
+    """This is the core function -- take a message and ingest the referred-to contents"""
 
     consignment_reference = message.get_consignment_reference()
     print(f"Ingester Start: Consignment reference {consignment_reference}")
@@ -539,3 +548,9 @@ def handler(event, context):
 
     print("Ingestion complete")
     return message.message
+
+
+@rollbar.lambda_function
+def handler(event, context):
+    for message in all_messages(event):
+        process_message(message)

ds-caselaw-ingester/tests.py

@@ -115,7 +115,9 @@ def test_handler_messages_v2(
         }
 
         message = v2_message_raw
-        event = {"Records": [{"Sns": {"Message": message}}]}
+        event = {
+            "Records": [{"Sns": {"Message": message}}, {"Sns": {"Message": message}}]
+        }
         lambda_function.handler(event=event, context=None)
 
         log = capsys.readouterr().out
@@ -127,13 +129,15 @@ def test_handler_messages_v2(
         assert "Ingestion complete" in log
         assert "auto_publish" not in log
         notify_update.assert_called()
+        assert notify_update.call_count == 2
         notify_new.assert_not_called()
         annotation.assert_called_with(
             ANY,
             automated=False,
             message="Updated document submitted by TDR user",
             payload=ANY,
         )
+        assert annotation.call_count == 2
 
     @patch("lambda_function.api_client", autospec=True)
     @patch("lambda_function.extract_metadata", autospec=True)
@@ -176,7 +180,9 @@ def test_handler_messages_s3(
         }
 
         message = s3_message_raw
-        event = {"Records": [{"Sns": {"Message": message}}]}
+        event = {
+            "Records": [{"Sns": {"Message": message}}, {"Sns": {"Message": message}}]
+        }
         lambda_function.handler(event=event, context=None)
 
         log = capsys.readouterr().out
@@ -188,6 +194,7 @@ def test_handler_messages_s3(
         assert "Ingestion complete" in log
         assert "auto_publish" in log
         apiclient.set_published.assert_called_with("failures/TDR-2020-FAR", True)
+        assert apiclient.set_published.call_count == 2
         notify_new.assert_not_called()
         notify_updated.assert_not_called()
         annotation.assert_called_with(
@@ -196,6 +203,7 @@ def test_handler_messages_s3(
             message="Updated document uploaded by Find Case Law",
             payload=ANY,
         )
+        assert annotation.call_count == 2
 
 
 class TestLambda: