diff --git a/modules/localstack/build.gradle b/modules/localstack/build.gradle index 97c5ad9ce35..5f4534b7756 100644 --- a/modules/localstack/build.gradle +++ b/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' } diff --git a/modules/localstack/src/main/java/org/testcontainers/containers/localstack/LocalStackContainer.java b/modules/localstack/src/main/java/org/testcontainers/containers/localstack/LocalStackContainer.java index b3651aea9e5..f0525af8176 100644 --- a/modules/localstack/src/main/java/org/testcontainers/containers/localstack/LocalStackContainer.java +++ b/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) { diff --git a/modules/localstack/src/test/java/org/testcontainers/containers/localstack/LocalstackContainerTest.java b/modules/localstack/src/test/java/org/testcontainers/containers/localstack/LocalstackContainerTest.java index ef0c9b728c0..a1c36078840 100644 --- a/modules/localstack/src/test/java/org/testcontainers/containers/localstack/LocalstackContainerTest.java +++ b/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().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 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 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(); + } + } }