[feature/PI-618-bulk_etl_e2e] add trigger tests

Author: jaklinger

src/etl/sds/tests/etl_test_utils/ask_s3.py

@@ -4,8 +4,9 @@
 
 from botocore.exceptions import ClientError
 from etl.clear_state_inputs import EMPTY_JSON_DATA, EMPTY_LDIF_DATA
-from etl_utils.constants import CHANGELOG_NUMBER, WorkerKey
+from etl_utils.constants import CHANGELOG_NUMBER, ETL_STATE_LOCK, WorkerKey
 from event.json import json_loads
+from mypy_boto3_dynamodb import DynamoDBClient
 from mypy_boto3_s3 import S3Client
 
 from etl.sds.tests.etl_test_utils.etl_state import EtlConfig
@@ -46,9 +47,31 @@ def ask_s3_prefix(
     return result
 
 
-def was_trigger_key_deleted(s3_client, bucket, etl_config: EtlConfig):
+def was_trigger_key_deleted(s3_client, etl_config: EtlConfig):
     return not ask_s3(
-        s3_client=s3_client, bucket=bucket, key=etl_config.initial_trigger_key
+        s3_client=s3_client,
+        bucket=etl_config.bucket,
+        key=etl_config.initial_trigger_key,
+    )
+
+
+def was_queue_history_created(
+    s3_client, etl_config: EtlConfig, expected_content: bytes
+):
+    return ask_s3_prefix(
+        s3_client=s3_client,
+        key_prefix=f"{etl_config.state_machine_history_key_prefix}/{etl_config.etl_type}",
+        question=lambda x: x == expected_content,
+    )
+
+
+def was_state_machine_history_created(
+    s3_client, etl_config: EtlConfig, expected_content: bytes
+):
+    return ask_s3_prefix(
+        s3_client=s3_client,
+        key_prefix=f"{etl_config.queue_history_key_prefix}/{etl_config.etl_type}",
+        question=lambda x: x == expected_content,
     )
 
 
@@ -61,27 +84,40 @@ def was_changelog_number_updated(s3_client, bucket, new_changelog_number):
     )
 
 
-def etl_state_is_clear(s3_client, bucket) -> bool:
-
-    extract_is_empty = ask_s3(
+def extract_is_empty(s3_client, bucket) -> bool:
+    return ask_s3(
         s3_client=s3_client,
         key=WorkerKey.EXTRACT,
         bucket=bucket,
         question=lambda x: x == EMPTY_LDIF_DATA,
     )
-    transform_is_empty = ask_s3(
+
+
+def transform_is_empty(s3_client, bucket) -> bool:
+    return ask_s3(
         s3_client=s3_client,
         key=WorkerKey.TRANSFORM,
         bucket=bucket,
         question=lambda x: x == EMPTY_JSON_DATA,
     )
-    load_is_empty = ask_s3(
+
+
+def load_is_empty(s3_client, bucket) -> bool:
+    return ask_s3(
         s3_client=s3_client,
         key=WorkerKey.LOAD,
         bucket=bucket,
         question=lambda x: x == EMPTY_JSON_DATA,
     )
-    print(  # noqa
-        f"{{extract, transform, load}} is empty == {{{extract_is_empty, transform_is_empty, load_is_empty}}}"
+
+
+def was_etl_state_lock_removed(s3_client, bucket) -> bool:
+    return ask_s3(
+        s3_client=s3_client,
+        key=ETL_STATE_LOCK,
+        bucket=bucket,
     )
-    return extract_is_empty and transform_is_empty and load_is_empty
+
+
+def database_isnt_empty(db_client: "DynamoDBClient", table_name: str):
+    return db_client.scan(TableName=table_name, Limit=1, Select="COUNT")["Count"] == 1

src/etl/sds/tests/etl_test_utils/etl_state.py

@@ -1,7 +1,13 @@
 from collections import deque
 from dataclasses import dataclass
 
-from etl_utils.constants import ETL_QUEUE_HISTORY, ETL_STATE_MACHINE_HISTORY, WorkerKey
+from etl_utils.constants import (
+    CHANGELOG_NUMBER,
+    ETL_QUEUE_HISTORY,
+    ETL_STATE_LOCK,
+    ETL_STATE_MACHINE_HISTORY,
+    WorkerKey,
+)
 from etl_utils.io import pkl_dumps_lz4
 from mypy_boto3_s3 import S3Client
 from sds.epr.bulk_create.bulk_load_fanout import FANOUT
@@ -18,16 +24,18 @@ class EtlConfig:
     initial_trigger_key: str
     queue_history_key_prefix: str
     state_machine_history_key_prefix: str
+    etl_type: str
     table_name: str
 
 
-def get_etl_config(input_filename: str) -> EtlConfig:
+def get_etl_config(input_filename: str, etl_type: str = "") -> EtlConfig:
     bulk_trigger_prefix = read_terraform_output("sds_etl.value.bulk_trigger_prefix")
     return EtlConfig(
         bucket=read_terraform_output("sds_etl.value.bucket"),
         initial_trigger_key=f"{bulk_trigger_prefix}/{input_filename}",
         queue_history_key_prefix=f"{ETL_QUEUE_HISTORY}/",
         state_machine_history_key_prefix=f"{ETL_STATE_MACHINE_HISTORY}/",
+        etl_type=etl_type,
         table_name=read_terraform_output("dynamodb_table_name.value"),
     )
 
@@ -74,3 +82,6 @@ def clear_etl_state(s3_client: S3Client, etl_config: EtlConfig):
         bucket=etl_config.bucket,
         key_prefix=etl_config.state_machine_history_key_prefix,
     )
+
+    s3_client.delete_object(Bucket=etl_config.bucket, Key=CHANGELOG_NUMBER)
+    s3_client.delete_object(Bucket=etl_config.bucket, Key=ETL_STATE_LOCK)

src/etl/sds/trigger/bulk/tests/_test_bulk_trigger.py

@@ -0,0 +1,127 @@
+import time
+from functools import partial
+
+import boto3
+import pytest
+from event.aws.client import dynamodb_client
+
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    database_isnt_empty as _database_isnt_empty,
+)
+from etl.sds.tests.etl_test_utils.ask_s3 import extract_is_empty as _extract_is_empty
+from etl.sds.tests.etl_test_utils.ask_s3 import load_is_empty as _load_is_empty
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    transform_is_empty as _transform_is_empty,
+)
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    was_changelog_number_updated as _was_changelog_number_updated,
+)
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    was_etl_state_lock_removed as _was_etl_state_lock_removed,
+)
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    was_queue_history_created as _was_queue_history_created,
+)
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    was_state_machine_history_created as _was_state_machine_history_created,
+)
+from etl.sds.tests.etl_test_utils.ask_s3 import (
+    was_trigger_key_deleted as _was_trigger_key_deleted,
+)
+from etl.sds.tests.etl_test_utils.etl_state import clear_etl_state, get_etl_config
+from etl.sds.worker.bulk.tests.test_bulk_e2e import PATH_TO_STAGE_DATA
+
+EXPECTED_CHANGELOG_NUMBER = 123
+
+
+def message(x):
+    print(x)  # noqa
+
+
+@pytest.mark.timeout(20)
+@pytest.mark.integration
+def test_bulk_trigger():
+    # Prerequisites
+    with open(PATH_TO_STAGE_DATA / "0.extract_input.ldif") as f:
+        input_data = f.read().encode()
+
+    etl_config = get_etl_config(f"{EXPECTED_CHANGELOG_NUMBER}.ldif", etl_type="bulk")
+    db_client = dynamodb_client()
+    s3_client = boto3.client("s3")
+    clear_etl_state(s3_client=s3_client, etl_config=etl_config)
+
+    # Define questions
+    was_trigger_key_deleted = partial(
+        _was_trigger_key_deleted, s3_client=s3_client, etl_config=etl_config
+    )
+    was_queue_history_created = partial(
+        _was_queue_history_created,
+        s3_client=s3_client,
+        etl_config=etl_config,
+        expected_content=input_data,
+    )
+    was_state_machine_history_created = partial(
+        _was_state_machine_history_created,
+        s3_client=s3_client,
+        etl_config=etl_config,
+        expected_content=input_data,
+    )
+    was_changelog_number_updated = partial(
+        _was_changelog_number_updated, s3_client=s3_client, bucket=etl_config.bucket
+    )
+    extract_is_empty = partial(
+        _extract_is_empty, s3_client=s3_client, bucket=etl_config.bucket
+    )
+    transform_is_empty = partial(
+        _transform_is_empty, s3_client=s3_client, bucket=etl_config.bucket
+    )
+    load_is_empty = partial(
+        _load_is_empty, s3_client=s3_client, bucket=etl_config.bucket
+    )
+    was_state_lock_removed = partial(
+        _was_etl_state_lock_removed, s3_client=s3_client, bucket=etl_config.bucket
+    )
+    database_isnt_empty = partial(
+        _database_isnt_empty, db_client=db_client, table_name=etl_config.table_name
+    )
+
+    # Trigger the bulk load
+    s3_client.put_object(
+        Bucket=etl_config.bucket, Key=etl_config.initial_trigger_key, Body=input_data
+    )
+
+    # Sign-off through the expected lifecycle of the bulk ETL
+    while not was_trigger_key_deleted():
+        time.sleep(5)
+    message("Trigger key deleted")
+
+    while not was_queue_history_created():
+        time.sleep(5)
+    message("Queue history created")
+
+    while not was_state_machine_history_created():
+        time.sleep(5)
+    message("State machine history created")
+
+    while not was_changelog_number_updated():
+        time.sleep(5)
+    message("Changelog number updated")
+
+    while not extract_is_empty():
+        time.sleep(5)
+    message("Extract's input data is now in empty state")
+
+    while not transform_is_empty():
+        time.sleep(5)
+    message("Transform's input data is now in empty state")
+
+    while not load_is_empty():
+        time.sleep(5)
+    message("Load's input data is now in empty state")
+
+    assert database_isnt_empty()
+    message("Database isn't empty")
+
+    while not was_state_lock_removed():
+        message("State lock has been removed")
+        time.sleep(5)