Skip to content

Commit

Permalink
Add LAMBDA_DOCKER_FLAGS with testcontainers labels (#8595)
Browse files Browse the repository at this point in the history
  • Loading branch information
dfangl authored May 6, 2024
1 parent 994b385 commit 2195610
Show file tree
Hide file tree
Showing 3 changed files with 138 additions and 0 deletions.
2 changes: 2 additions & 0 deletions modules/localstack/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -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'
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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.
Expand Down Expand Up @@ -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();
Expand All @@ -185,6 +215,7 @@ protected void configure() {
}

exposePorts();
configureLambdaContainerLabels();
}

private void resolveHostname(String envVar) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;

Expand Down Expand Up @@ -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();
}
}
}

0 comments on commit 2195610

Please sign in to comment.