start_lambda_api_integ_base.py

import shutil
import uuid
from typing import Optional, Dict, List
from unittest import TestCase, skipIf
import threading
from subprocess import Popen, PIPE
import os
import logging
from pathlib import Path

import boto3
from botocore import UNSIGNED
from botocore.config import Config
from docker.errors import APIError
from psutil import NoSuchProcess

from samcli.local.docker.utils import get_validated_container_client
from tests.integration.local.common_utils import random_port, InvalidAddressException, wait_for_local_process

from tests.testing_utils import (
    SKIP_DOCKER_TESTS,
    SKIP_DOCKER_MESSAGE,
    run_command,
    kill_process,
    get_sam_command,
)

LOG = logging.getLogger(__name__)


@skipIf(SKIP_DOCKER_TESTS, SKIP_DOCKER_MESSAGE)
class StartLambdaIntegBaseClass(TestCase):
    template: Optional[str] = None
    container_mode: Optional[str] = None
    container_host_interface: Optional[str] = None
    parameter_overrides: Optional[Dict[str, str]] = None
    binary_data_file: Optional[str] = None
    integration_dir = str(Path(__file__).resolve().parents[2])
    invoke_image: Optional[List] = None
    hook_name: Optional[str] = None
    terraform_plan_file: Optional[str] = None
    beta_features: Optional[bool] = None
    collect_start_lambda_process_output: bool = False
    function_logical_ids: Optional[List[str]] = None

    build_before_invoke = False
    build_overrides: Optional[Dict[str, str]] = None

    working_dir: Optional[str] = None

    @classmethod
    def setUpClass(cls):
        # This is the directory for tests/integration which will be used to file the testdata
        # files for integ tests
        cls.template = cls.integration_dir + cls.template_path
        cls.working_dir = str(Path(cls.template).resolve().parents[0])
        cls.env_var_path = cls.integration_dir + "/testdata/invoke/vars.json"

        if cls.build_before_invoke:
            cls.build()

        # Create the Docker client with automatic container selection
        cls.docker_client = get_validated_container_client()

        # Only remove containers with SAM CLI labels to avoid interfering with other processes
        try:
            sam_containers = cls.docker_client.containers.list(
                all=True, filters={"label": "sam.cli.container.type=lambda"}
            )
            for container in sam_containers:
                try:
                    container.remove(force=True)
                    LOG.info("Removed existing SAM CLI container %s", container.short_id)
                except APIError as ex:
                    LOG.error("Failed to remove existing SAM CLI container %s", container.short_id, exc_info=ex)
        except Exception as ex:
            LOG.error("Failed to clean up existing SAM CLI containers", exc_info=ex)

        cls.start_lambda_with_retry()

    @classmethod
    def build(cls):
        command = get_sam_command()

        command_list = [command, "build"]
        if cls.build_overrides:
            overrides_arg = " ".join(
                ["ParameterKey={},ParameterValue={}".format(key, value) for key, value in cls.build_overrides.items()]
            )
            command_list += ["--parameter-overrides", overrides_arg]
        if cls.hook_name:
            command_list += ["--hook-name", cls.hook_name]
        run_command(command_list, cwd=cls.working_dir)

    @classmethod
    def start_lambda_with_retry(cls, retries=3, input=None, env=None):
        retry_count = 0
        while retry_count < retries:
            cls.port = str(random_port())
            try:
                cls.start_lambda(input=input, env=env)
            except InvalidAddressException:
                retry_count += 1
                continue
            break

        if retry_count == retries:
            raise ValueError("Ran out of retries attempting to start lambda")

    @classmethod
    def get_start_lambda_command(
        cls,
        port=None,
        template_path=None,
        env_var_path=None,
        container_mode=None,
        container_host_interface=None,
        parameter_overrides=None,
        invoke_image=None,
        hook_name=None,
        beta_features=None,
        terraform_plan_file=None,
        function_logical_ids=None,
    ):
        command_list = [get_sam_command(), "local", "start-lambda"]

        # Add function names as positional arguments first
        if function_logical_ids:
            command_list.extend(function_logical_ids)

        if port:
            command_list += ["-p", port]

        if template_path:
            command_list += ["-t", template_path]

        if env_var_path:
            command_list += ["--env-vars", env_var_path]

        if container_mode:
            command_list += ["--warm-containers", container_mode]

        if container_host_interface:
            command_list += ["--container-host-interface", container_host_interface]

        if parameter_overrides:
            command_list += ["--parameter-overrides", cls._make_parameter_override_arg(parameter_overrides)]

        if invoke_image:
            for image in invoke_image:
                command_list += ["--invoke-image", image]

        if hook_name:
            command_list += ["--hook-name", hook_name]

        if beta_features is not None:
            command_list += ["--beta-features" if beta_features else "--no-beta-features"]

        if terraform_plan_file:
            command_list += ["--terraform-plan-file", terraform_plan_file]

        return command_list

    @classmethod
    def start_lambda(cls, wait_time=5, input=None, env=None):
        command_list = cls.get_start_lambda_command(
            port=cls.port,
            template_path=cls.template,
            env_var_path=cls.env_var_path,
            container_mode=cls.container_mode,
            container_host_interface=cls.container_host_interface,
            parameter_overrides=cls.parameter_overrides,
            invoke_image=cls.invoke_image,
            hook_name=cls.hook_name,
            beta_features=cls.beta_features,
            terraform_plan_file=cls.terraform_plan_file,
            function_logical_ids=cls.function_logical_ids,
        )

        # Container labels are no longer needed - container IDs are parsed from output
        cls.start_lambda_process = Popen(command_list, stderr=PIPE, stdin=PIPE, env=env, cwd=cls.working_dir)
        cls.start_lambda_process_output = ""

        if input:
            cls.start_lambda_process.stdin.write(input)
            cls.start_lambda_process.stdin.close()

        cls.start_lambda_process_error = wait_for_local_process(
            cls.start_lambda_process, cls.port, cls.collect_start_lambda_process_output
        )

        cls.stop_reading_thread = False

        def read_sub_process_stderr():
            while not cls.stop_reading_thread:
                line = cls.start_lambda_process.stderr.readline()
                if line:
                    line_str = line.decode("utf-8").strip()
                    cls.start_lambda_process_output += line_str + "\n"

        cls.read_threading = threading.Thread(target=read_sub_process_stderr, daemon=True)
        cls.read_threading.start()

    @classmethod
    def _make_parameter_override_arg(self, overrides):
        return " ".join(["ParameterKey={},ParameterValue={}".format(key, value) for key, value in overrides.items()])

    @classmethod
    def tearDownClass(cls):
        # After all the tests run, we need to kill the start_lambda process.
        cls.stop_reading_thread = True
        try:
            kill_process(cls.start_lambda_process)
        except NoSuchProcess:
            LOG.info("Process has already been terminated")

        # Clean up any remaining SAM CLI containers
        try:
            docker_client = get_validated_container_client()
            # Only remove containers with SAM CLI labels to avoid interfering with other processes
            sam_containers = docker_client.containers.list(all=True, filters={"label": "sam.cli.container.type=lambda"})
            for container in sam_containers:
                try:
                    container.remove(force=True)
                    LOG.info("Removed SAM CLI container %s", container.short_id)
                except APIError as ex:
                    LOG.error("Failed to remove SAM CLI container %s", container.short_id, exc_info=ex)
        except Exception as ex:
            LOG.error("Failed to clean up SAM CLI containers", exc_info=ex)

    def get_local_lambda_client(self):
        url = "http://127.0.0.1:{}".format(self.port)
        return boto3.client(
            "lambda",
            endpoint_url=url,
            region_name="us-east-1",
            use_ssl=False,
            verify=False,
            config=Config(signature_version=UNSIGNED, read_timeout=120, retries={"max_attempts": 0}),
        )


class WatchWarmContainersIntegBaseClass(StartLambdaIntegBaseClass):
    temp_path: Optional[str] = None
    template_path: Optional[str] = None
    code_path: Optional[str] = None
    docker_file_path: Optional[str] = None

    @classmethod
    def setUpClass(cls):
        cls.temp_path = str(uuid.uuid4()).replace("-", "")[:10]
        working_dir = str(Path(cls.integration_dir).resolve().joinpath(cls.temp_path))
        if Path(working_dir).resolve().exists():
            shutil.rmtree(working_dir, ignore_errors=True)
        os.mkdir(working_dir)
        os.mkdir(Path(cls.integration_dir).resolve().joinpath(cls.temp_path).joinpath("dir"))
        cls.template_path = f"/{cls.temp_path}/template.yaml"
        cls.code_path = f"/{cls.temp_path}/main.py"
        cls.code_path2 = f"/{cls.temp_path}/dir/main2.py"
        cls.docker_file_path = f"/{cls.temp_path}/Dockerfile"
        cls.docker_file_path2 = f"/{cls.temp_path}/Dockerfile2"

        if cls.template_content:
            cls._write_file_content(cls.template_path, cls.template_content)

        if cls.code_content:
            cls._write_file_content(cls.code_path, cls.code_content)

        if cls.docker_file_content:
            cls._write_file_content(cls.docker_file_path, cls.docker_file_content)

        super().setUpClass()

    @classmethod
    def _write_file_content(cls, path, content):
        with open(cls.integration_dir + path, "w") as f:
            f.write(content)

    @classmethod
    def tearDownClass(cls):
        super().tearDownClass()
        working_dir = str(Path(cls.integration_dir).resolve().joinpath(cls.temp_path))
        if Path(working_dir).resolve().exists():
            shutil.rmtree(working_dir, ignore_errors=True)