"""Shared fixtures and utilities for Step Function integration tests.""" import json import os import random import shutil import tempfile import time import zipfile from datetime import datetime, timezone from typing import Any, Callable import boto3 import pytest import requests from common.connectors.jwt_connector import create_test_session from config import QA_BASE_URL, TEST_VENDOR_UUID from jwtauth.testing import JwtAuthSecretsManager, SecretLookupInfo # Enable jwtauth testing plugin for JWT generation pytest_plugins = ["jwtauth.testing.pytest_plugin"] @pytest.fixture def sfn_client(): """Create Step Functions client.""" return boto3.client("stepfunctions", region_name=os.getenv("AWS_DEFAULT_REGION", "us-east-1")) @pytest.fixture def s3_client(): """Create S3 client.""" return boto3.client("s3", region_name=os.getenv("AWS_DEFAULT_REGION", "us-east-1")) @pytest.fixture def aws_account_id(): """Get AWS account ID.""" return boto3.client("sts").get_caller_identity()["Account"] @pytest.fixture def aws_region(): """Get AWS region.""" return os.getenv("AWS_DEFAULT_REGION", "us-east-1") @pytest.fixture(scope="session") def jwt( generate_bearer_token: Callable[..., str], jwtauth_secrets_manager: JwtAuthSecretsManager, ) -> str: """Generate JWT via AWS Secrets Manager + Auth0 OAuth flow.""" return generate_bearer_token( get_user_creds_args=SecretLookupInfo( environment="qa", service_name="ows-product-staging-integration-test", secret_name="OWS_PRODUCT_STAGING_INTEGRATION_TEST_USER_CREDENTIALS", ), get_auth0_creds_args=SecretLookupInfo( environment="qa", service_name="ows-product-staging-integration-test", secret_name="OWS_PRODUCT_STAGING_INTEGRATION_TEST_APP_AUTH0_CREDENTIALS", ), secrets_manager=jwtauth_secrets_manager, ) @pytest.fixture def test_bulk_session(jwt: str) -> dict[str, Any]: """Create bulk session for testing using JWT authentication. Returns dict with 'id' field containing bulk_session_id. """ return create_test_session( vendor_uuid=TEST_VENDOR_UUID, jwt_token=jwt, base_url=QA_BASE_URL ) @pytest.fixture def identity_uuid(): """Return a fixed identity UUID for testing purposes.""" return "10436b38-5e11-472d-b6a4-bf1ee2b1b438" def generate_upc_check_digit(first_11_digits: str) -> str: """Calculate UPC check digit using the standard algorithm.""" odd_sum = sum(int(first_11_digits[i]) for i in range(0, 11, 2)) even_sum = sum(int(first_11_digits[i]) for i in range(1, 11, 2)) total = (odd_sum * 3) + even_sum check_digit = (10 - (total % 10)) % 10 return str(check_digit) def generate_unique_upc() -> str: """Generate a unique 12-digit UPC with valid check digit.""" timestamp = str(int(time.time()))[-8:] random_part = str(random.randint(100, 999)) first_11 = timestamp + random_part check_digit = generate_upc_check_digit(first_11) return first_11 + check_digit def generate_unique_product_code(prefix: str = "TST") -> str: """Generate a unique product code in ABC-123 format.""" timestamp = int(time.time()) random_suffix = random.randint(100, 999) return f"{prefix}-{timestamp % 10000}{random_suffix}" def update_excel_with_unique_codes( source_file_path: str, old_project_code: str, new_project_code: str, old_upc: str, new_upc: str, old_manufacturers_upc: str, new_manufacturers_upc: str, old_product_code: str, new_product_code: str, ) -> bytes: """Update Excel file with new unique codes and return the modified file as bytes. Args: source_file_path: Path to the source Excel file old_project_code: Old project code to replace (e.g., "TST-7980") new_project_code: New project code (e.g., "TST-123456789") old_upc: Old UPC to replace (e.g., "679779803435") new_upc: New UPC (e.g., "123456789012") old_manufacturers_upc: Old manufacturers UPC to replace new_manufacturers_upc: New manufacturers UPC old_product_code: Old product code to replace (e.g., "TEST-7980") new_product_code: New product code (e.g., "TEST-123456789") Returns: bytes: Modified Excel file as bytes """ # Create temporary directory temp_dir = tempfile.mkdtemp() try: # Extract Excel file with zipfile.ZipFile(source_file_path, 'r') as z: z.extractall(temp_dir) # Update sharedStrings.xml (where most text is stored) shared_strings_path = os.path.join(temp_dir, 'xl/sharedStrings.xml') if os.path.exists(shared_strings_path): with open(shared_strings_path, 'r', encoding='utf-8') as f: content = f.read() content = content.replace(old_project_code, new_project_code) content = content.replace(old_upc, new_upc) content = content.replace(old_manufacturers_upc, new_manufacturers_upc) content = content.replace(old_product_code, new_product_code) with open(shared_strings_path, 'w', encoding='utf-8') as f: f.write(content) # Update worksheet files (for inline values like UPC) for i in range(1, 20): sheet_path = os.path.join(temp_dir, f'xl/worksheets/sheet{i}.xml') if os.path.exists(sheet_path): with open(sheet_path, 'r', encoding='utf-8') as f: content = f.read() original_content = content content = content.replace(old_upc, new_upc) content = content.replace(old_manufacturers_upc, new_manufacturers_upc) if content != original_content: with open(sheet_path, 'w', encoding='utf-8') as f: f.write(content) # Create new Excel file in memory output_path = os.path.join(temp_dir, 'output.xlsx') with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as z: for root, dirs, files in os.walk(temp_dir): for file in files: if file == 'output.xlsx': continue file_path = os.path.join(root, file) arcname = os.path.relpath(file_path, temp_dir) z.write(file_path, arcname) # Read the output file as bytes with open(output_path, 'rb') as f: return f.read() finally: # Cleanup shutil.rmtree(temp_dir, ignore_errors=True) def wait_for_execution(sfn_client, execution_arn, timeout=300): """Wait for step function execution to complete. Args: sfn_client: Boto3 Step Functions client execution_arn: ARN of the execution to wait for timeout: Maximum time to wait in seconds (default 5 minutes) Returns: dict: Execution response with status """ start_time = time.time() while time.time() - start_time < timeout: response = sfn_client.describe_execution(executionArn=execution_arn) status = response["status"] if status in ["SUCCEEDED", "FAILED", "TIMED_OUT", "ABORTED"]: return response time.sleep(5) # Poll every 5 seconds raise TimeoutError(f"Execution did not complete within {timeout} seconds") def wait_for_execution_to_appear(sfn_client, state_machine_arn, start_time, timeout=60): """Wait for a step function execution to appear after S3 upload. Args: sfn_client: Boto3 Step Functions client state_machine_arn: ARN of the state machine start_time: Timestamp when the S3 upload occurred timeout: Maximum time to wait in seconds (default 60s) Returns: str: Execution ARN of the found execution """ wait_start = time.time() while time.time() - wait_start < timeout: # List recent executions - check all statuses, not just RUNNING # The step function might complete quickly and already be in SUCCEEDED/FAILED state for status in ["RUNNING", "SUCCEEDED", "FAILED", "TIMED_OUT", "ABORTED"]: response = sfn_client.list_executions( stateMachineArn=state_machine_arn, statusFilter=status, maxResults=10 ) # Look for execution started after our upload for execution in response.get("executions", []): exec_start_time = execution["startDate"] if exec_start_time >= start_time: return execution["executionArn"] time.sleep(2) # Poll every 2 seconds raise TimeoutError(f"No execution of {state_machine_arn} appeared within {timeout} seconds after S3 upload") def upload_test_file_to_s3(s3_client, bucket_name, file_key, content, metadata=None): """Upload a test file to S3. Args: s3_client: Boto3 S3 client bucket_name: Name of the S3 bucket file_key: S3 key for the file content: File content (will be JSON-encoded if dict/list, or bytes/file path) metadata: Optional dict of S3 metadata to attach to the object Returns: datetime: Timestamp when the upload occurred (timezone-aware UTC) """ upload_time = datetime.now(timezone.utc) # Handle different content types if isinstance(content, (dict, list)): body = json.dumps(content) content_type = "application/json" elif isinstance(content, bytes): body = content content_type = "application/octet-stream" elif isinstance(content, str) and os.path.exists(content): # If content is a file path, read the file with open(content, "rb") as f: body = f.read() # Determine content type based on file extension if content.endswith(".xlsx"): content_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" elif content.endswith(".json"): content_type = "application/json" else: content_type = "application/octet-stream" else: body = content content_type = "text/plain" put_params = { "Bucket": bucket_name, "Key": file_key, "Body": body, "ContentType": content_type } if metadata: put_params["Metadata"] = metadata s3_client.put_object(**put_params) return upload_time def cleanup_s3_file(s3_client, bucket_name, file_key): """Delete a test file from S3. Args: s3_client: Boto3 S3 client bucket_name: Name of the S3 bucket file_key: S3 key for the file to delete """ try: s3_client.delete_object(Bucket=bucket_name, Key=file_key) except Exception: pass # Best effort cleanup def get_state_machine_arn(state_machine_name, account_id, region): """Construct the ARN for a Step Function state machine. Args: state_machine_name: Name of the state machine account_id: AWS account ID region: AWS region Returns: str: Full ARN of the state machine """ return f"arn:aws:states:{region}:{account_id}:stateMachine:{state_machine_name}" def verify_execution_success(sfn_client, execution_arn, execution_result): """Verify that a step function execution succeeded. Args: sfn_client: Boto3 Step Functions client execution_arn: ARN of the execution execution_result: Execution response from describe_execution Raises: AssertionError: If execution did not succeed """ status = execution_result["status"] if status != "SUCCEEDED": # Get execution history for debugging any failure history = sfn_client.get_execution_history(executionArn=execution_arn) error_message = ( f"Step function execution failed with status: {status}\n" f"Execution history: {json.dumps(history, indent=2, default=str)}" ) raise AssertionError(error_message) assert status == "SUCCEEDED", f"Expected SUCCEEDED status, got {status}" def register_asset_via_api( bulk_session_id: str, original_filename: str, jwt: str, base_url: str = QA_BASE_URL ) -> tuple[str, str]: """Register an asset through the API and return S3 details. Args: bulk_session_id: ID of the bulk session original_filename: Original filename of the asset jwt: JWT for authentication base_url: Base URL for OWS Product Staging API Returns: tuple: (s3_filename, presigned_url) where: - s3_filename: The S3 key generated by the API (e.g., "assets/{uuid}") - presigned_url: Presigned URL for uploading the file Raises: requests.exceptions.RequestException: If the API call fails """ url = f"{base_url.rstrip('/')}/assets/upload" headers = {"Authorization": f"Bearer {jwt}"} payload = { "bulk_session_id": bulk_session_id, "assets": [ { "original_filename": original_filename, "parts": 1 # Single-part upload for small test files } ] } response = requests.post(url, json=payload, headers=headers, timeout=30) response.raise_for_status() response_data = response.json() # Extract the S3 filename and presigned URL from the first asset if response_data.get("assets") and len(response_data["assets"]) > 0: asset = response_data["assets"][0] s3_filename = asset["s3_filename"] # Get the first presigned URL (for single-part upload) if asset.get("presigned_urls") and len(asset["presigned_urls"]) > 0: presigned_url = asset["presigned_urls"][0]["url"] return s3_filename, presigned_url raise ValueError("API did not return presigned URLs for the asset") raise ValueError("API did not return asset data") def wait_for_metadata_validation(bulk_session_id: str, jwt: str, timeout: int = 60, base_url: str = QA_BASE_URL) -> None: """Wait for metadata validation to complete.""" url = f"{base_url.rstrip('/')}/bulk-session/{bulk_session_id}" headers = {"Authorization": f"Bearer {jwt}"} start_time = time.time() while time.time() - start_time < timeout: try: response = requests.get(url, headers=headers, timeout=10) response.raise_for_status() session_data = response.json() metadata_status = session_data.get("metadata_status") if metadata_status == "valid": return # If status is error or failed, raise immediately if metadata_status in ["error", "failed", "invalid"]: raise RuntimeError(f"Metadata validation failed with status: {metadata_status}") except requests.exceptions.RequestException: pass time.sleep(2) # Poll every 2 seconds raise TimeoutError(f"Metadata did not become valid within {timeout} seconds") def upload_metadata_with_unique_codes( bulk_session_id: str, original_filename: str, source_file_path: str, jwt: str, project_code_prefix: str = "TST", product_code_prefix: str = "TEST", base_url: str = QA_BASE_URL, ) -> tuple[str, dict[str, str], datetime, str]: """Upload metadata file with dynamically generated unique codes. Args: bulk_session_id: ID of the bulk session original_filename: Original filename to register (e.g., "bulk_upload_file.xlsx") source_file_path: Path to the source Excel file jwt: JWT for authentication project_code_prefix: Prefix for project code (default "TST") product_code_prefix: Prefix for product code (default "TEST") base_url: Base URL for OWS Product Staging API Returns: tuple: (s3_filename, codes_dict, upload_time, etag) where: - s3_filename: The S3 key for the uploaded file - codes_dict: Dictionary with keys: project_code, upc, manufacturers_upc, product_code - upload_time: Upload timestamp - etag: Upload ETag """ # Generate unique codes new_upc = generate_unique_upc() new_manufacturers_upc = new_upc[:-1] # 11 digits new_project_code = generate_unique_product_code(project_code_prefix) new_product_code = generate_unique_product_code(product_code_prefix) codes_dict = { "project_code": new_project_code, "upc": new_upc, "manufacturers_upc": new_manufacturers_upc, "product_code": new_product_code, } # Determine old values based on filename if "ingest" in original_filename.lower(): old_project_code = "ING-1191" old_upc = "679811906438" old_manufacturers_upc = "67981190643" old_product_code = "INGST-1191" else: old_project_code = "TST-7980" old_upc = "679779803435" old_manufacturers_upc = "67977980343" old_product_code = "TEST-7980" # Update Excel file with new codes updated_file_bytes = update_excel_with_unique_codes( source_file_path, old_project_code, new_project_code, old_upc, new_upc, old_manufacturers_upc, new_manufacturers_upc, old_product_code, new_product_code, ) # Register with API s3_filename, presigned_urls = register_metadata_via_api( bulk_session_id, original_filename, jwt, base_url ) # Upload the modified file upload_time = datetime.now(timezone.utc) headers = {} if not ("uploadId=" in presigned_urls[0] and "partNumber=" in presigned_urls[0]): headers["Content-Type"] = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" response = requests.put(presigned_urls[0], data=updated_file_bytes, headers=headers, timeout=60) response.raise_for_status() etag = response.headers.get('ETag', '').strip('"') return s3_filename, codes_dict, upload_time, etag def upload_via_presigned_url(presigned_url: str, file_path: str, metadata: dict[str, str] | None = None) -> tuple[datetime, str]: """Upload file using presigned URL.""" upload_time = datetime.now(timezone.utc) with open(file_path, "rb") as f: file_data = f.read() # Prepare headers headers = {} is_multipart_upload = "uploadId=" in presigned_url and "partNumber=" in presigned_url if not is_multipart_upload: # Only add Content-Type for regular (non-multipart) uploads if file_path.endswith(".mp3"): headers["Content-Type"] = "audio/mpeg" elif file_path.endswith(".m4a"): headers["Content-Type"] = "audio/mp4" elif file_path.endswith(".wav"): headers["Content-Type"] = "audio/wav" elif file_path.endswith(".xlsx"): headers["Content-Type"] = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" elif file_path.endswith(".jpg") or file_path.endswith(".jpeg"): headers["Content-Type"] = "image/jpeg" elif file_path.endswith(".tiff") or file_path.endswith(".tif"): headers["Content-Type"] = "image/tiff" else: headers["Content-Type"] = "application/octet-stream" if metadata: for key, value in metadata.items(): headers[f"x-amz-meta-{key}"] = value # Upload to presigned URL response = requests.put(presigned_url, data=file_data, headers=headers, timeout=60) response.raise_for_status() # Extract ETag from response headers (strip quotes if present) etag = response.headers.get('ETag', '').strip('"') return upload_time, etag def register_metadata_via_api( bulk_session_id: str, original_filename: str, jwt: str, base_url: str = QA_BASE_URL ) -> tuple[str, list[str]]: """Register metadata upload and get presigned URLs via API. This mimics the real-world client flow for metadata uploads: 1. Client calls POST /metadata/upload to register metadata upload 2. API creates bulk_session_metadata_file record with embedded S3 metadata 3. API returns S3 filename and presigned URL(s) for upload 4. Client uploads using the presigned URLs 5. Client calls PATCH /metadata/upload/{s3_filename} to complete the upload Args: bulk_session_id: ID of the bulk session original_filename: Original name of the metadata file jwt: JWT for authentication base_url: Base URL for OWS Product Staging API Returns: tuple: (s3_filename, [presigned_urls]) for the metadata upload Raises: requests.exceptions.RequestException: If the API call fails """ url = f"{base_url.rstrip('/')}/metadata/upload" headers = {"Authorization": f"Bearer {jwt}"} payload = { "bulk_session_id": bulk_session_id, "original_filename": original_filename, } response = requests.post(url, json=payload, headers=headers, timeout=30) response.raise_for_status() metadata_file = response.json() s3_filename = metadata_file["s3_filename"] # Now get presigned URLs for upload presigned_url_response = requests.get( f"{base_url.rstrip('/')}/metadata/upload/{s3_filename}", params={"parts": 1}, # Single-part upload for test files headers=headers, timeout=30, ) presigned_url_response.raise_for_status() presigned_urls_data = presigned_url_response.json() presigned_urls = [url["url"] for url in presigned_urls_data.get("presigned_urls", [])] if not presigned_urls: raise ValueError("API did not return presigned URLs for metadata") return s3_filename, presigned_urls def complete_metadata_upload( s3_filename: str, etags: list[str], jwt: str, base_url: str = QA_BASE_URL ) -> None: """Complete metadata multipart upload.""" url = f"{base_url.rstrip('/')}/metadata/upload/{s3_filename}" headers = {"Authorization": f"Bearer {jwt}"} # Build parts list with ETags parts = [{"part_number": i + 1, "etag": etag} for i, etag in enumerate(etags)] payload = {"parts": parts} response = requests.patch(url, json=payload, headers=headers, timeout=30) response.raise_for_status() def complete_asset_upload( s3_filename: str, etags: list[str], jwt: str, base_url: str = QA_BASE_URL ) -> None: """Complete asset multipart upload with retry logic for transient errors.""" url = f"{base_url.rstrip('/')}/assets/upload/{s3_filename}" headers = {"Authorization": f"Bearer {jwt}"} # Build parts list with ETags parts = [{"part_number": i + 1, "etag": etag} for i, etag in enumerate(etags)] payload = {"parts": parts} # Retry logic for transient 502/503 errors max_retries = 3 retry_delay = 2 for attempt in range(max_retries): try: response = requests.patch(url, json=payload, headers=headers, timeout=30) response.raise_for_status() return # Success except requests.exceptions.HTTPError as e: if e.response.status_code in [502, 503] and attempt < max_retries - 1: time.sleep(retry_delay) retry_delay *= 2 # Exponential backoff else: raise def wait_for_audio_encoding_completion( bulk_session_id: str, jwt: str, timeout: int = 180, base_url: str = QA_BASE_URL ) -> None: """Wait for all audio assets to finish encoding in ows-assets. This function polls the bulk session to verify that all assets are ready for ingestion by checking the asset_status field. Args: bulk_session_id: ID of the bulk session jwt: JWT for authentication timeout: Maximum time to wait in seconds (default 3 minutes) base_url: Base URL for OWS Product Staging API Raises: TimeoutError: If audio encoding doesn't complete within timeout RuntimeError: If audio encoding fails """ session_url = f"{base_url.rstrip('/')}/bulk-session/{bulk_session_id}" headers = {"Authorization": f"Bearer {jwt}"} start_time = time.time() while time.time() - start_time < timeout: try: response = requests.get(session_url, headers=headers, timeout=10) response.raise_for_status() session_data = response.json() asset_status = session_data.get("asset_status") # Check if asset_status is "complete" (assets are ready for ingestion) if asset_status == "complete": return # Check for error states if asset_status == "error": raise RuntimeError(f"Asset processing failed for bulk session {bulk_session_id}") except requests.exceptions.RequestException: pass time.sleep(5) # Poll every 5 seconds raise TimeoutError(f"Asset processing did not complete within {timeout} seconds") def get_bulk_session_state(bulk_session_id: str, jwt: str, base_url: str = QA_BASE_URL) -> dict: """Get the current state of a bulk session. Args: bulk_session_id: ID of the bulk session jwt: JWT for authentication base_url: Base URL for OWS Product Staging API Returns: dict: Bulk session data Raises: requests.exceptions.RequestException: If the API call fails """ url = f"{base_url.rstrip('/')}/bulk-session/{bulk_session_id}" headers = {"Authorization": f"Bearer {jwt}"} response = requests.get(url, headers=headers, timeout=10) response.raise_for_status() return response.json() def trigger_ingestion(bulk_session_id: str, jwt: str, base_url: str = QA_BASE_URL) -> dict: """Trigger ingestion for a bulk session. This sends a POST request to create a bulk session ingestion, which will trigger the ingestion step function to process all products in the session. Args: bulk_session_id: ID of the bulk session to ingest jwt: JWT for authentication base_url: Base URL for OWS Product Staging API Returns: dict: Response from the API containing ingestion details including bulk_session_ingestion_id Raises: requests.exceptions.HTTPError: If the API returns an error (e.g., 422 if assets not ready) """ url = f"{base_url.rstrip('/')}/bulk-session/{bulk_session_id}" headers = {"Authorization": f"Bearer {jwt}"} # Send POST request with required fields payload = { "assets_required": True, "submit": True, "send_notifications": False } response = requests.post(url, json=payload, headers=headers, timeout=30) response.raise_for_status() return response.json()