Function workflow_copy_to_s3

test/aws-localstack/mzcompose.py:244–304

(c: Composition)

def workflow_copy_to_s3(c: Composition) -> None:
    with c.override(
        Materialized(
            depends_on=["localstack"],
            environment_extra=[
                f"AWS_ENDPOINT_URL={AWS_ENDPOINT_URL_MZ}",
                f"AWS_ACCESS_KEY_ID={AWS_ACCESS_KEY_ID}",
                f"AWS_SECRET_ACCESS_KEY={AWS_SECRET_ACCESS_KEY}",
            ],
        )
    ):
        c.up("localstack", "materialized")
        localhost_aws_endpoint_url = f"http://localhost:{c.port('localstack', 4566)}"
        s3_client = boto3.client(
            "s3",
            endpoint_url=localhost_aws_endpoint_url,
            region_name=DEFAULT_CLOUD_REGION,
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
        )
        bucket_name = "copy-to-s3"
        s3_client.create_bucket(Bucket=bucket_name)
        path_prefix = str(uuid.uuid4())
        c.run_testdrive_files(
            f"--var=endpoint={AWS_ENDPOINT_URL_MZ}",
            f"--var=access-key={AWS_ACCESS_KEY_ID}",
            f"--var=secret-key={AWS_SECRET_ACCESS_KEY}",
            f"--var=s3-prefix={bucket_name}/{path_prefix}",
            f"--var=region={DEFAULT_CLOUD_REGION}",
            "--default-timeout=300s",
            "copy-to-s3/copy-to-s3.td",
        )

        def validate_upload(upload, expected_output_set):
            assert len(upload["Contents"]) > 0
            output_lines = set()
            for obj in upload["Contents"]:
                assert obj["Key"].endswith(".csv")
                key = obj["Key"]
                object_response = s3_client.get_object(Bucket=bucket_name, Key=key)
                body = object_response["Body"].read().decode("utf-8")
                output_lines.update(body.splitlines())
            assert output_lines == expected_output_set

        # asserting the uploaded files
        date = c.sql_query("SELECT TO_CHAR(now(), 'YYYY-MM-DD')")[0][0]
        expected_output = set(map(lambda x: str(x), range(10)))
        first_upload = s3_client.list_objects_v2(
            Bucket=bucket_name, Prefix=f"{path_prefix}/1/{date}/"
        )
        validate_upload(first_upload, expected_output)

        second_upload = s3_client.list_objects_v2(
            Bucket=bucket_name, Prefix=f"{path_prefix}/2/"
        )
        validate_upload(second_upload, expected_output)

        third_upload = s3_client.list_objects_v2(