cdk/test/test_cdk_post_deploy.py

"""Unit tests for cdk_post_deploy.py (boto3 helpers, no live AWS)."""

import json
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch

CDK_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(CDK_DIR))

import cdk_post_deploy as post


def test_container_definitions_with_named_port_adds_mapping_to_first_container():
    containers = [{"name": "app", "image": "nginx"}]
    updated = post._container_definitions_with_named_port(
        containers,
        port_name="port-7860",
        container_port=7860,
    )
    assert updated[0]["portMappings"] == [
        {"name": "port-7860", "containerPort": 7860, "protocol": "tcp"}
    ]


def test_container_definitions_with_named_port_names_existing_mapping():
    containers = [
        {
            "name": "app",
            "image": "nginx",
            "portMappings": [{"containerPort": 7860, "protocol": "tcp"}],
        }
    ]
    updated = post._container_definitions_with_named_port(
        containers,
        port_name="port-7860",
        container_port=7860,
    )
    assert updated[0]["portMappings"][0]["name"] == "port-7860"


def test_resolve_service_task_definition_arn_from_describe_services():
    mock_ecs = MagicMock()
    mock_ecs.describe_services.return_value = {
        "services": [
            {
                "serviceArn": "arn:aws:ecs:eu-west-2:123:service/cluster/app",
                "taskDefinition": "arn:aws:ecs:eu-west-2:123:task-definition/app:1",
            }
        ]
    }

    with patch("cdk_post_deploy.boto3.client", return_value=mock_ecs):
        arn = post.resolve_service_task_definition_arn("cluster", "app")

    assert arn == "arn:aws:ecs:eu-west-2:123:task-definition/app:1"
    mock_ecs.describe_express_gateway_service.assert_not_called()


def test_resolve_service_task_definition_arn_from_express_service_revision():
    mock_ecs = MagicMock()
    mock_ecs.describe_services.return_value = {
        "services": [
            {
                "serviceArn": "arn:aws:ecs:eu-west-2:123:service/cluster/express-app",
            }
        ]
    }
    mock_ecs.describe_express_gateway_service.return_value = {
        "service": {
            "activeConfigurations": [
                {
                    "serviceRevisionArn": "arn:aws:ecs:eu-west-2:123:service-revision/rev/1"
                }
            ]
        }
    }
    mock_ecs.describe_service_revisions.return_value = {
        "serviceRevisions": [
            {
                "taskDefinition": "arn:aws:ecs:eu-west-2:123:task-definition/express:3",
            }
        ]
    }

    with patch("cdk_post_deploy.boto3.client", return_value=mock_ecs):
        arn = post.resolve_service_task_definition_arn("cluster", "express-app")

    assert arn == "arn:aws:ecs:eu-west-2:123:task-definition/express:3"
    mock_ecs.describe_express_gateway_service.assert_called_once()
    mock_ecs.describe_service_revisions.assert_called_once_with(
        serviceRevisionArns=["arn:aws:ecs:eu-west-2:123:service-revision/rev/1"]
    )


def test_start_express_gateway_service_updates_scaling_target():
    mock_ecs = MagicMock()
    mock_ecs.get_paginator.return_value.paginate.return_value = [
        {
            "serviceArns": [
                "arn:aws:ecs:eu-west-2:123456789012:service/my-cluster/my-express"
            ]
        }
    ]

    with patch("cdk_post_deploy.boto3.client", return_value=mock_ecs):
        result = post.start_express_gateway_service("my-cluster", "my-express")

    assert result["statusCode"] == 200
    mock_ecs.update_express_gateway_service.assert_called_once_with(
        serviceArn="arn:aws:ecs:eu-west-2:123456789012:service/my-cluster/my-express",
        scalingTarget=post.EXPRESS_GATEWAY_ACTIVE_SCALING_TARGET,
    )


def test_cognito_https_callback_urls():
    assert post.cognito_https_callback_urls(
        "https://abc123.eu-west-2.elb.amazonaws.com"
    ) == [
        "https://abc123.eu-west-2.elb.amazonaws.com",
        "https://abc123.eu-west-2.elb.amazonaws.com/oauth2/idpresponse",
    ]
    assert post.cognito_https_callback_urls("app.example.com")[0].startswith("https://")


def test_update_user_pool_client_callback_urls_preserves_oauth_settings():
    mock_cognito = MagicMock()
    mock_cognito.describe_user_pool_client.return_value = {
        "UserPoolClient": {
            "ClientName": "app-client",
            "CallbackURLs": ["https://old.example.com"],
            "AllowedOAuthFlows": ["code"],
            "AllowedOAuthScopes": ["openid", "email", "profile"],
            "AllowedOAuthFlowsUserPoolClient": True,
            "SupportedIdentityProviders": ["COGNITO"],
            "ExplicitAuthFlows": ["ALLOW_REFRESH_TOKEN_AUTH"],
        }
    }

    with patch("cdk_post_deploy.boto3.client", return_value=mock_cognito):
        post.update_user_pool_client_callback_urls(
            "pool-1",
            "client-1",
            [
                "https://new.example.com",
                "https://new.example.com/oauth2/idpresponse",
            ],
            aws_region="eu-west-2",
        )

    mock_cognito.update_user_pool_client.assert_called_once()
    kwargs = mock_cognito.update_user_pool_client.call_args.kwargs
    assert kwargs["UserPoolId"] == "pool-1"
    assert kwargs["ClientId"] == "client-1"
    assert kwargs["CallbackURLs"] == [
        "https://new.example.com",
        "https://new.example.com/oauth2/idpresponse",
    ]
    assert kwargs["AllowedOAuthFlows"] == ["code"]
    assert kwargs["AllowedOAuthScopes"] == ["openid", "email", "profile"]


def test_apply_cognito_alb_callback_fixup_skips_when_already_correct():
    mock_cognito = MagicMock()
    mock_cognito.describe_user_pool_client.return_value = {
        "UserPoolClient": {
            "CallbackURLs": post.cognito_https_callback_urls("https://app.example.com"),
        }
    }

    with patch("cdk_post_deploy.boto3.client", return_value=mock_cognito):
        changed = post.apply_cognito_alb_callback_fixup(
            user_pool_id="pool-1",
            client_id="client-1",
            redirect_base="https://app.example.com",
        )

    assert changed is False
    mock_cognito.update_user_pool_client.assert_not_called()


def test_target_group_arn_from_ecs_register_event():
    message = (
        "(service my-svc) registered 1 targets in "
        "(target-group arn:aws:elasticloadbalancing:eu-west-2:123:"
        "targetgroup/ecs-gateway-tg-abc/def)"
    )
    assert (
        post.target_group_arn_from_ecs_register_event(message)
        == "arn:aws:elasticloadbalancing:eu-west-2:123:targetgroup/ecs-gateway-tg-abc/def"
    )


def test_resolve_express_service_target_group_arn_waits_for_registration_event():
    mock_ecs = MagicMock()
    mock_ecs.describe_services.side_effect = [
        {
            "services": [
                {
                    "events": [],
                    "runningCount": 0,
                    "desiredCount": 1,
                }
            ]
        },
        {
            "services": [
                {
                    "events": [
                        {
                            "message": (
                                "(service Demo-Summarisation-ECSService) registered 1 targets in "
                                "(target-group arn:aws:elasticloadbalancing:eu-west-2:123:"
                                "targetgroup/ecs-gateway-tg-abc/def)"
                            )
                        }
                    ],
                    "runningCount": 1,
                    "desiredCount": 1,
                }
            ]
        },
    ]

    with (
        patch("cdk_post_deploy.boto3.client", return_value=mock_ecs),
        patch("cdk_post_deploy.time.sleep"),
        patch("cdk_post_deploy.time.monotonic", side_effect=[0, 0, 600]),
    ):
        arn = post.resolve_express_service_target_group_arn(
            "cluster",
            "Demo-Summarisation-ECSService",
            max_wait_seconds=600,
            poll_interval_seconds=15,
        )

    assert arn == (
        "arn:aws:elasticloadbalancing:eu-west-2:123:targetgroup/ecs-gateway-tg-abc/def"
    )
    assert mock_ecs.describe_services.call_count == 2


def test_resolve_express_service_target_group_arn_from_task_ips():
    mock_ecs = MagicMock()
    mock_ecs.describe_services.return_value = {
        "services": [
            {
                "events": [],
                "runningCount": 1,
                "desiredCount": 1,
            }
        ]
    }
    mock_ecs.list_tasks.return_value = {"taskArns": ["arn:task/1"]}
    mock_ecs.describe_tasks.return_value = {
        "tasks": [
            {
                "attachments": [
                    {"details": [{"name": "privateIPv4Address", "value": "10.0.1.42"}]}
                ]
            }
        ]
    }
    mock_elbv2 = MagicMock()
    mock_elbv2.describe_target_groups.return_value = {
        "TargetGroups": [{"TargetGroupArn": "arn:tg/main"}]
    }
    mock_elbv2.describe_target_health.return_value = {
        "TargetHealthDescriptions": [{"Target": {"Id": "10.0.1.42"}}]
    }

    def client_factory(service_name, **_kwargs):
        return mock_elbv2 if service_name == "elbv2" else mock_ecs

    with (
        patch("cdk_post_deploy.boto3.client", side_effect=client_factory),
        patch("cdk_post_deploy.time.monotonic", return_value=0),
    ):
        arn = post.resolve_express_service_target_group_arn(
            "cluster",
            "Demo-Summarisation-ECSService",
            load_balancer_arn="arn:lb/express",
            max_wait_seconds=600,
            poll_interval_seconds=15,
        )

    assert arn == "arn:tg/main"


def test_listener_actions_with_target_group_replaces_forward_arn():
    actions = [
        {"Type": "authenticate-cognito", "Order": 1},
        {
            "Type": "forward",
            "Order": 2,
            "TargetGroupArn": "arn:old",
            "ForwardConfig": {"TargetGroups": [{"TargetGroupArn": "arn:old"}]},
        },
    ]
    updated = post.listener_actions_with_target_group(actions, "arn:new")
    forward = next(action for action in updated if action["Type"] == "forward")
    assert forward["TargetGroupArn"] == "arn:new"
    assert forward["ForwardConfig"]["TargetGroups"][0]["TargetGroupArn"] == "arn:new"


def test_cognito_secret_payload_matches():
    desired = {
        "SUMMARISATION_USER_POOL_ID": "eu-west-2_AAAA",
        "SUMMARISATION_CLIENT_ID": "client",
        "SUMMARISATION_CLIENT_SECRET": "secret",
    }
    current = json.dumps(
        {
            "SUMMARISATION_USER_POOL_ID": "eu-west-2_OLD",
            "SUMMARISATION_CLIENT_ID": "client",
            "SUMMARISATION_CLIENT_SECRET": "secret",
        }
    )
    assert post.cognito_secret_payload_matches(current, desired) is False
    assert post.cognito_secret_payload_matches(json.dumps(desired), desired) is True


def test_listener_rule_has_cognito_auth():
    assert post.listener_rule_has_cognito_auth(
        [{"Type": "authenticate-cognito"}, {"Type": "forward"}]
    )
    assert not post.listener_rule_has_cognito_auth([{"Type": "forward"}])


def test_print_express_mode_next_steps(capsys, monkeypatch):
    def fake_get_stack_output(stack_name, output_key, region):
        outputs = {
            "ExpressServiceEndpoint": "main.example.ecs.eu-west-2.on.aws",
            "PiExpressEndpoint": "pi.example.ecs.eu-west-2.on.aws",
        }
        return outputs.get(output_key)

    monkeypatch.setattr(post, "get_stack_output", fake_get_stack_output)
    post.print_express_mode_next_steps(
        {
            "AWS_REGION": "eu-west-2",
            "ENABLE_PI_AGENT_EXPRESS_SERVICE": "True",
        }
    )
    out = capsys.readouterr().out
    assert "Wait 10 minutes for app deployment to finish." in out
    assert "Cognito authorisation" not in out
    assert "sign in at the Pi agent URL" in out
    assert "https://main.example.ecs.eu-west-2.on.aws" in out
    assert "https://pi.example.ecs.eu-west-2.on.aws/" in out
    assert "pi_agent.env" not in out


def test_print_headless_deployment_next_steps(capsys):
    post.print_headless_deployment_next_steps(
        {
            "AWS_REGION": "eu-west-2",
            "S3_OUTPUT_BUCKET_NAME": "my-output-bucket",
            "S3_LOG_CONFIG_BUCKET_NAME": "my-log-bucket",
            "S3_BATCH_INPUT_PREFIX": "input/",
            "S3_BATCH_ENV_PREFIX": "input/config/",
            "S3_BATCH_DEFAULT_PARAMS_KEY": "general-config/app_defaults.env",
            "S3_BATCH_LAMBDA_FUNCTION_NAME": "Headless-Summarisation-S3BatchEcsTrigger",
            "ECS_LOG_GROUP_NAME": "/ecs/headless-summarisation-ecsservice-logs",
        }
    )
    out = capsys.readouterr().out
    assert (
        "Upload a consultation spreadsheet (.xlsx) to s3://my-output-bucket/input/"
        in out
    )
    assert "example_headless_env_file.env" in out
    assert "s3://my-output-bucket/input/config/" in out
    assert "Headless-Summarisation-S3BatchEcsTrigger" in out
    assert "s3://my-output-bucket/output/<session-folder>/" in out
    assert "s3://my-log-bucket/general-config/app_defaults.env" in out
    assert "tools/config.py" in out


def test_seed_headless_batch_s3_layout_creates_prefixes(tmp_path):
    from botocore.exceptions import ClientError

    example = tmp_path / "example_headless_env_file.env"
    example.write_text("DIRECT_MODE_TASK=redact\n", encoding="utf-8")
    defaults = tmp_path / "app_defaults.env"
    defaults.write_text("RUN_DIRECT_MODE=1\nS3_OUTPUTS_BUCKET=\n", encoding="utf-8")

    missing = ClientError(
        {"Error": {"Code": "404", "Message": "Not Found"}},
        "HeadObject",
    )
    s3 = MagicMock()
    s3.head_object.side_effect = missing

    with patch("cdk_post_deploy.boto3.client", return_value=s3):
        post.seed_headless_batch_s3_layout(
            "my-bucket",
            example_env_local_path=str(example),
            log_bucket="my-log-bucket",
            app_defaults_local_path=str(defaults),
            s3_outputs_bucket_name="my-bucket",
            aws_region="eu-west-2",
        )

    put_calls = s3.put_object.call_args_list
    assert len(put_calls) == 4
    keys_by_bucket = [(call.kwargs["Bucket"], call.kwargs["Key"]) for call in put_calls]
    assert ("my-bucket", "input/") in keys_by_bucket
    assert ("my-bucket", "input/config/") in keys_by_bucket
    assert ("my-bucket", "input/config/example_headless_env_file.env") in keys_by_bucket
    assert ("my-log-bucket", "general-config/app_defaults.env") in keys_by_bucket