import calendar from collections.abc import Generator from datetime import date from typing import Any import pytest from aws_testing_utils.step_function_handler import StepFunctionHandler from test_fixtures.mysql import MySQLConnection from test_fixtures.snowflake import SnowflakeConnection from tests.helpers import discover_single_pk, placeholders _ORIGINATING_VENDOR_ID = 6971 _DESTINATION_VENDOR_ID = 7123 _TEST_IDENTITY_ID = "testtest-test-4est-8est-testtesttest" _SNOWFLAKE_SCHEMA = "QA" _DIM_RELEASE_HISTORY = f"FACTS.{_SNOWFLAKE_SCHEMA}.DIM_RELEASE_HISTORY" _DIM_RELEASE = f"FACTS.{_SNOWFLAKE_SCHEMA}.DIM_RELEASE" _STATEMENT_PERIOD = ( f"ORCHARD_APP_REPORTING_V2.{_SNOWFLAKE_SCHEMA}_ROYALTY_ACCOUNTING_ROYALTY_ACCOUNTING.STATEMENT_PERIOD" ) @pytest.fixture def sfn_full_job( art_relations_db: MySQLConnection, snowflake_db: SnowflakeConnection, ) -> Generator[dict[str, Any], None, None]: """Composite fixture satisfying all lambdas in the state machine. Combines the seeding requirements of EnsureArtists, ExecuteContentTransfer, and UpdateDimTables into a single job: - originating_artist_id set (EnsureArtists reads this for the bulk-ensure call) - product_transfer_history rows with source_artist_id; destination_artist_id is left NULL so EnsureArtists populates it at runtime before ExecuteContentTransfer runs - revenue_cutoff_date = last day of the month 2 months ago (UpdateDimTables) - only releases that exist in DIM_RELEASE are included (UpdateDimTables INSERT path) Captures project/release/video/release_artist pre-state so teardown can restore them safely even if the SFN fails mid-flight and ExecuteContentTransfer left things partially mutated. """ # 1. revenue_cutoff_date = end of the month 2 months ago. # The lambda looks up STATEMENT_PERIOD for MONTH(cutoff) directly, so we fetch that # period to derive expected_start_period for the assertion. today = date.today() cutoff_month = today.month - 2 cutoff_year = today.year if cutoff_month <= 0: cutoff_month += 12 cutoff_year -= 1 revenue_cutoff_date = date(cutoff_year, cutoff_month, calendar.monthrange(cutoff_year, cutoff_month)[1]).isoformat() period = snowflake_db.fetchone( f"SELECT STATEMENT_PERIOD_ID FROM {_STATEMENT_PERIOD} WHERE STATEMENT_YEAR = %s AND STATEMENT_MONTH = %s", (cutoff_year, cutoff_month), ) assert period, ( f"No STATEMENT_PERIOD found for {cutoff_year}-{cutoff_month:02d}; " f"UpdateDimTables assertion will be inaccurate — check that QA has a period for that month." ) expected_start_period = int(period["STATEMENT_PERIOD_ID"]) + 1 # 2. Project with 2-5 releases and artist_id set (ExecuteContentTransfer requirement) project = art_relations_db.fetchone( """ SELECT project_id, vendor_id, subaccount_id, artist_id FROM project WHERE vendor_id = %s AND artist_id IS NOT NULL AND (SELECT COUNT(*) FROM releases WHERE project_id = project.project_id) BETWEEN 2 AND 5 LIMIT 1 """, (_ORIGINATING_VENDOR_ID,), ) assert project, ( f"No transferable project found for vendor_id {_ORIGINATING_VENDOR_ID} (need 2–5 releases and artist_id set)" ) # 3. All releases for this project all_releases = art_relations_db.fetchall( "SELECT release_id, artist_id, subaccount_id FROM releases WHERE project_id = %s", (project["project_id"],), ) assert all_releases, f"Project {project['project_id']} unexpectedly has no releases" all_release_ids = [r["release_id"] for r in all_releases] # 4. Keep only releases that exist in Snowflake DIM_RELEASE (UpdateDimTables INSERT path) dim_rows = snowflake_db.fetchall( f"SELECT DISTINCT PRODUCT_ID FROM {_DIM_RELEASE} WHERE PRODUCT_ID IN ({placeholders(all_release_ids)})", tuple(all_release_ids), ) snowflake_ids = {int(r["PRODUCT_ID"]) for r in dim_rows} releases = [r for r in all_releases if r["release_id"] in snowflake_ids][:2] assert releases, ( f"None of project {project['project_id']}'s releases exist in {_DIM_RELEASE}; " f"cannot drive the UpdateDimTables INSERT path." ) release_ids = [r["release_id"] for r in releases] # 5. Capture pre-state for teardown across ALL releases in the project. # ExecuteContentTransfer moves the entire project — not just the Snowflake-filtered # subset — so every release, video, and release_artist row must be restored. product_video_pk = discover_single_pk(art_relations_db, "product_video") video_originals = art_relations_db.fetchall( f"SELECT {product_video_pk} AS pk, release_id, primary_artist_id " f"FROM product_video WHERE release_id IN ({placeholders(all_release_ids)})", tuple(all_release_ids), ) release_artist_originals = art_relations_db.fetchall( f"SELECT release_artist_id, release_id, artist_info_id " f"FROM release_artist WHERE release_id IN ({placeholders(all_release_ids)})", tuple(all_release_ids), ) # 6. Seed transfer job created_by = _TEST_IDENTITY_ID art_relations_db.execute( """ INSERT INTO project_transfer_job (project_id, originating_vendor_id, originating_artist_id, destination_vendor_id, revenue_cutoff_date, created_by_identity_id) VALUES (%s, %s, %s, %s, %s, %s) """, ( project["project_id"], _ORIGINATING_VENDOR_ID, project["artist_id"], _DESTINATION_VENDOR_ID, revenue_cutoff_date, created_by, ), ) job = art_relations_db.fetchone("SELECT * FROM project_transfer_job WHERE job_id = LAST_INSERT_ID()") assert job is not None # 7. Seed product_transfer_history — destination_artist_id intentionally NULL; # EnsureArtists populates it before ExecuteContentTransfer runs. for release in releases: art_relations_db.execute( "INSERT INTO product_transfer_history (job_id, release_id, source_artist_id) VALUES (%s, %s, %s)", (job["job_id"], release["release_id"], release["artist_id"]), ) yield { "job": job, "project": project, "releases": releases, "release_ids": release_ids, "expected_start_period": expected_start_period, } # Teardown: restore mutated rows first, then remove seed data. # UPDATE is safe even if ExecuteContentTransfer didn't run (no-op when values unchanged). art_relations_db.execute( "UPDATE project SET vendor_id = %s, subaccount_id = %s, artist_id = %s WHERE project_id = %s", (project["vendor_id"], project["subaccount_id"], project["artist_id"], project["project_id"]), ) for r in all_releases: art_relations_db.execute( "UPDATE releases SET artist_id = %s, subaccount_id = %s WHERE release_id = %s", (r["artist_id"], r["subaccount_id"], r["release_id"]), ) for v in video_originals: art_relations_db.execute( f"UPDATE product_video SET primary_artist_id = %s WHERE {product_video_pk} = %s", (v["primary_artist_id"], v["pk"]), ) for ra in release_artist_originals: art_relations_db.execute( "UPDATE release_artist SET artist_info_id = %s WHERE release_artist_id = %s", (ra["artist_info_id"], ra["release_artist_id"]), ) art_relations_db.execute("DELETE FROM product_transfer_history WHERE job_id = %s", (job["job_id"],)) art_relations_db.execute("DELETE FROM project_transfer_job WHERE job_id = %s", (job["job_id"],)) pytestmark = pytest.mark.lambda_name("state-machine") def test_happy_path_succeeds_and_all_datastores_updated( art_relations_db: MySQLConnection, snowflake_db: SnowflakeConnection, sfn_full_job: dict[str, Any], sfn_handler: StepFunctionHandler, ) -> None: """End-to-end: the full state machine runs to SUCCEEDED and every datastore is updated. Covers all lambdas in sequence: - EnsureArtists: populates destination_artist_id in product_transfer_history - ExecuteContentTransfer: moves project + releases to destination vendor - UpdateDimTables: closes open DIM_RELEASE_HISTORY rows and inserts new ones - EditAttachments: no-op for QA test jobs (no staged terms); still must not fail - FinalizeJob: sets project_transfer_job.status = COMPLETED Assertions probe three distinct datastores so a partial failure (e.g. only FinalizeJob ran) would be caught even if SFN reports SUCCEEDED. """ job_id = sfn_full_job["job"]["job_id"] # Start the execution with assert_success=False so we can capture the ARN and # stop the execution before fixture teardown runs if the poll times out or fails. sfn_handler.execute({"job_id": job_id}, assert_success=False) execution_arn = sfn_handler.get_last_execution() try: sfn_handler.assert_execution_success(execution_arn, max_wait=300) except Exception: try: sfn_handler.client.stop_execution(executionArn=execution_arn) except Exception: pass raise # FinalizeJob: job must be COMPLETED job_row = art_relations_db.fetchone( "SELECT status, transfer_completed_on FROM project_transfer_job WHERE job_id = %s", (job_id,), ) assert job_row is not None assert job_row["status"] == "COMPLETED", ( f"Expected job {job_id} status COMPLETED after SFN SUCCEEDED, got {job_row['status']!r}" ) assert job_row["transfer_completed_on"] is not None, ( f"transfer_completed_on is NULL for job {job_id} despite COMPLETED status" ) # ExecuteContentTransfer: project must have moved to destination vendor project_row = art_relations_db.fetchone( "SELECT vendor_id FROM project WHERE project_id = %s", (sfn_full_job["project"]["project_id"],), ) assert project_row is not None assert project_row["vendor_id"] == _DESTINATION_VENDOR_ID, ( f"project.vendor_id was not updated to {_DESTINATION_VENDOR_ID} for job {job_id}; " f"got {project_row['vendor_id']!r}" ) # UpdateDimTables: each release must have an open DIM_RELEASE_HISTORY row at the destination expected_start_period = sfn_full_job["expected_start_period"] for release_id in sfn_full_job["release_ids"]: open_rows = snowflake_db.fetchall( f"SELECT LABELID, START_STATEMENT_PERIOD_ID " f"FROM {_DIM_RELEASE_HISTORY} " f"WHERE PRODUCT_ID = %s AND END_DATE_INCLUSIVE IS NULL", (release_id,), ) assert open_rows, ( f"No open DIM_RELEASE_HISTORY row for release {release_id} after SFN run (UpdateDimTables did not insert)" ) for row in open_rows: assert int(row["LABELID"]) == _DESTINATION_VENDOR_ID, ( f"Release {release_id}: open DIM_RELEASE_HISTORY row has LABELID {row['LABELID']}, " f"expected destination vendor {_DESTINATION_VENDOR_ID}" ) assert int(row["START_STATEMENT_PERIOD_ID"]) == expected_start_period, ( f"Release {release_id}: open row has START_STATEMENT_PERIOD_ID " f"{row['START_STATEMENT_PERIOD_ID']}, expected {expected_start_period} " f"(a stale open row means the close step didn't run)" )