Add LAMBDA_DOCKER_FLAGS with testcontainers labels (#8595)

Author: dfangl

modules/localstack/build.gradle

@@ -7,6 +7,8 @@ dependencies {
     testImplementation 'com.amazonaws:aws-java-sdk-s3'
     testImplementation 'com.amazonaws:aws-java-sdk-sqs'
     testImplementation 'com.amazonaws:aws-java-sdk-logs'
+    testImplementation 'com.amazonaws:aws-java-sdk-lambda'
+    testImplementation 'com.amazonaws:aws-java-sdk-core'
     testImplementation 'software.amazon.awssdk:s3:2.23.9'
     testImplementation 'org.assertj:assertj-core:3.25.2'
 }

modules/localstack/src/main/java/org/testcontainers/containers/localstack/LocalStackContainer.java

@@ -10,6 +10,7 @@
 import org.testcontainers.containers.wait.strategy.Wait;
 import org.testcontainers.utility.ComparableVersion;
 import org.testcontainers.utility.DockerImageName;
+import org.testcontainers.utility.ResourceReaper;
 
 import java.net.InetAddress;
 import java.net.URI;
@@ -19,6 +20,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * Testcontainers implementation for LocalStack.
@@ -163,6 +165,34 @@ private static boolean shouldRunInLegacyMode(String version) {
         return true;
     }
 
+    /**
+     * Provides a docker argument string including all default labels set on testcontainer containers
+     * @return Argument string in the format `-l key1=value1 -l key2=value2`
+     */
+    private static String internalMarkerLabels() {
+        return Stream
+            .concat(
+                DockerClientFactory.DEFAULT_LABELS.entrySet().stream(),
+                ResourceReaper.instance().getLabels().entrySet().stream()
+            )
+            .map(entry -> String.format("-l %s=%s", entry.getKey(), entry.getValue()))
+            .collect(Collectors.joining(" "));
+    }
+
+    /**
+     * Configure the LocalStack container to include the default testcontainer labels on all spawned lambda containers
+     * Necessary to properly clean up lambda containers even if the LocalStack container is killed before it gets the
+     * chance.
+     */
+    private void configureLambdaContainerLabels() {
+        String lambdaDockerFlags = internalMarkerLabels();
+        String existingLambdaDockerFlags = getEnvMap().get("LAMBDA_DOCKER_FLAGS");
+        if (existingLambdaDockerFlags != null) {
+            lambdaDockerFlags = existingLambdaDockerFlags + " " + lambdaDockerFlags;
+        }
+        withEnv("LAMBDA_DOCKER_FLAGS", lambdaDockerFlags);
+    }
+
     @Override
     protected void configure() {
         super.configure();
@@ -185,6 +215,7 @@ protected void configure() {
         }
 
         exposePorts();
+        configureLambdaContainerLabels();
     }
 
     private void resolveHostname(String envVar) {

modules/localstack/src/test/java/org/testcontainers/containers/localstack/LocalstackContainerTest.java

@@ -8,6 +8,15 @@
 import com.amazonaws.services.kms.model.CreateKeyRequest;
 import com.amazonaws.services.kms.model.CreateKeyResult;
 import com.amazonaws.services.kms.model.Tag;
+import com.amazonaws.services.lambda.AWSLambda;
+import com.amazonaws.services.lambda.AWSLambdaClientBuilder;
+import com.amazonaws.services.lambda.model.CreateFunctionRequest;
+import com.amazonaws.services.lambda.model.CreateFunctionResult;
+import com.amazonaws.services.lambda.model.FunctionCode;
+import com.amazonaws.services.lambda.model.GetFunctionRequest;
+import com.amazonaws.services.lambda.model.InvokeRequest;
+import com.amazonaws.services.lambda.model.InvokeResult;
+import com.amazonaws.services.lambda.model.Runtime;
 import com.amazonaws.services.logs.AWSLogs;
 import com.amazonaws.services.logs.AWSLogsClientBuilder;
 import com.amazonaws.services.logs.model.CreateLogGroupRequest;
@@ -20,12 +29,15 @@
 import com.amazonaws.services.sqs.AmazonSQS;
 import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
 import com.amazonaws.services.sqs.model.CreateQueueResult;
+import com.amazonaws.waiters.WaiterParameters;
+import com.github.dockerjava.api.DockerClient;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.io.IOUtils;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.experimental.runners.Enclosed;
 import org.junit.runner.RunWith;
+import org.testcontainers.DockerClientFactory;
 import org.testcontainers.containers.Container;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.Network;
@@ -36,14 +48,21 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.s3.S3Client;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.net.URL;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.Date;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -505,4 +524,90 @@ public void shouldBeAccessibleWithCredentials() throws IOException {
             assertThat(content).as("The object can be retrieved").isEqualTo("baz");
         }
     }
+
+    public static class LambdaContainerLabels {
+
+        @ClassRule
+        public static LocalStackContainer localstack = new LocalStackContainer(
+            LocalstackTestImages.LOCALSTACK_2_3_IMAGE
+        );
+
+        private static byte[] createLambdaHandlerZipFile() throws IOException {
+            StringBuilder sb = new StringBuilder();
+            sb.append("def handler(event, context):\n");
+            sb.append("    return event");
+
+            ByteArrayOutputStream byteOutput = new ByteArrayOutputStream();
+            ZipOutputStream out = new ZipOutputStream(byteOutput);
+            ZipEntry e = new ZipEntry("handler.py");
+            out.putNextEntry(e);
+
+            byte[] data = sb.toString().getBytes();
+            out.write(data, 0, data.length);
+            out.closeEntry();
+            out.close();
+            return byteOutput.toByteArray();
+        }
+
+        @Test
+        public void shouldLabelLambdaContainers() throws IOException {
+            AWSLambda lambda = AWSLambdaClientBuilder
+                .standard()
+                .withEndpointConfiguration(
+                    new AwsClientBuilder.EndpointConfiguration(
+                        localstack.getEndpoint().toString(),
+                        localstack.getRegion()
+                    )
+                )
+                .withCredentials(
+                    new AWSStaticCredentialsProvider(
+                        new BasicAWSCredentials(localstack.getAccessKey(), localstack.getSecretKey())
+                    )
+                )
+                .build();
+
+            // create function
+            byte[] handlerFile = createLambdaHandlerZipFile();
+            CreateFunctionRequest createFunctionRequest = new CreateFunctionRequest()
+                .withFunctionName("test-function")
+                .withRuntime(Runtime.Python311)
+                .withHandler("handler.handler")
+                .withRole("arn:aws:iam::000000000000:role/test-role")
+                .withCode(new FunctionCode().withZipFile(ByteBuffer.wrap(handlerFile)));
+            CreateFunctionResult createFunctionResult = lambda.createFunction(createFunctionRequest);
+            GetFunctionRequest getFunctionRequest = new GetFunctionRequest()
+                .withFunctionName(createFunctionResult.getFunctionName());
+            lambda
+                .waiters()
+                .functionActiveV2()
+                .run(new WaiterParameters<GetFunctionRequest>().withRequest(getFunctionRequest));
+
+            // invoke function once
+            String payload = "{\"test\": \"payload\"}";
+            InvokeRequest invokeRequest = new InvokeRequest()
+                .withFunctionName(createFunctionResult.getFunctionName())
+                .withPayload(payload);
+            InvokeResult invokeResult = lambda.invoke(invokeRequest);
+            assertThat(StandardCharsets.UTF_8.decode(invokeResult.getPayload()).toString())
+                .as("Invoke result not matching expected output")
+                .isEqualTo(payload);
+
+            // assert that the spawned lambda containers has the testcontainers labels set
+            DockerClient dockerClient = DockerClientFactory.instance().client();
+            Collection<String> nameFilter = Collections.singleton(localstack.getContainerName().replace("_", "-"));
+            com.github.dockerjava.api.model.Container lambdaContainer = dockerClient
+                .listContainersCmd()
+                .withNameFilter(nameFilter)
+                .exec()
+                .stream()
+                .findFirst()
+                .orElse(null);
+            assertThat(lambdaContainer).as("Lambda container not found").isNotNull();
+            Map<String, String> labels = lambdaContainer.getLabels();
+            assertThat(labels.get("org.testcontainers")).as("TestContainers label not present").isEqualTo("true");
+            assertThat(labels.get("org.testcontainers.sessionId"))
+                .as("TestContainers session id not present")
+                .isNotNull();
+        }
+    }
 }