[SVLS-8146] add LMI integ tests

.gitignore

@@ -26,5 +26,7 @@ local_tests/serverless-init/logs.txt
 bottlecap/target
 bottlecap/proptest-regressions
 
+integration-tests/cdk.context.json
+
 .gitlab/pipeline*
 /CLAUDE.md

.gitlab/datasources/test-suites.yaml

@@ -2,3 +2,4 @@ test_suites:
   - name: base
   - name: otlp
   - name: snapstart
+  - name: lmi

integration-tests/README.md

@@ -13,89 +13,22 @@ The general flow is:
 For simplicity, integration tests are set up to only test against ARM runtimes and 4 runtimes (Python, Node, Java, and Dotnet).
 
 
-## Test Suites
-
-### Base Tests
-
-The base test suite provides basic functionality tests across all supported Lambda runtimes. These tests verify core extension functionality without additional instrumentation.
-
-**What it tests:**
-- Extension can collect and forward logs to Datadog
-- Extension generates and sends traces with proper span structure
-- Extension detects cold starts correctly
-
-**Test Coverage:**
-- Lambda invocation succeeds (200 status code)
-- "Hello world!" log message is sent to Datadog
-- One trace is sent to Datadog
-- `aws.lambda` span exists with correct properties including `cold_start: 'true'`
-- `aws.lambda.cold_start` span is created
-- `aws.lambda.load` span is created for Python and Node
-
-### OTLP Tests
-
-The OTLP test suite verifies OpenTelemetry Protocol (OTLP) integration with the Datadog Lambda Extension. These tests use Lambda functions instrumented with OpenTelemetry SDKs to ensure telemetry data flows correctly through the extension to Datadog.
-
-**What it tests:**
-- Lambda functions instrumented with OpenTelemetry SDKs can invoke successfully
-- Traces are properly sent to Datadog via OTLP
-- Spans contain correct structure and attributes
-
-**Test Coverage:**
-- Lambda invocation succeeds (200 status code)
-- At least one trace is sent to Datadog
-- Trace contains valid spans with proper structure
-
-## CI/CD Pipeline Structure
-
 
-### Pipeline Flow
-
-```
-                                    ┌→ deploy-base → test-base → cleanup-base ┐
-publish layer → build lambdas ─────┤                                          ├→ cleanup-layer
-                                    └→ deploy-otlp → test-otlp → cleanup-otlp ┘
-```
-
-### Test Suite Lifecycle
-
-Each test suite (base, otlp, etc.) follows this lifecycle:
-
-1. **Deploy**: Deploys only the stacks for that suite
-   - Pattern: `cdk deploy "integ-${IDENTIFIER}-${TEST_SUITE}"`
-   - Example: `integ-abc123-base` deploys all base test stacks
-
-2. **Test**: Runs only the tests for that suite
-   - Command: `jest tests/${TEST_SUITE}.test.ts`
-   - Example: `jest tests/base.test.ts`
-
-3. **Cleanup**: Removes only the stacks for that suite
-   - Runs with `when: always` to ensure cleanup on failure
-   - Pattern: Deletes all stacks matching `integ-${IDENTIFIER}-${TEST_SUITE}`
+## Test Suites
+### Overview
+* Gitlab runs test suites in parallel.
+* For each test suite, Gitlab will:
+  * Deploy the associated stack
+  * Run the test suite
+  * Destroy the associated stack
+* You can download the test suite run from the Gitlab page. 
 
 ### Adding a New Test Suite
 
-To add a new test suite to the parallel execution:
-
 1. **Create test file**: `tests/<suite-name>.test.ts`
 2. **Create CDK stacks**: `lib/stacks/<suite-name>.ts`
 3. **Register stacks**: Add to `bin/app.ts`
-4. **Update pipeline**: Add suite name to `.gitlab/templates/pipeline.yaml.tpl`:
-   ```yaml
-   parallel:
-     matrix:
-       - TEST_SUITE: [base, otlp, <suite-name>]
-   ```
-
-## Guidelines
-
-### Naming
-* An `identifier` is used to differentiate the different stacks. For local development, the identifier is automatically set using the command `whoami` and parses the user's first name. For gitlab pipelines, the identifier is the git commit short sha.
-* Stacks should be named `integ-<identifier>-<stack name>`
-* Lambda functions should be named `<stack-id>-<function name>`
-
-### Stacks
-* Stacks automatically get destroyed at the end of the gitlab integration tests step. Stack should be setup to not retain resources. A helper function `createLogGroup` exists with `removalPolicy: cdk.RemovalPolicy.DESTROY`. 
+4. **Update pipeline**: Add suite name to `.gitlab/datasources/test-suites.yaml`:
 
 
 ## Local Development
@@ -140,7 +73,6 @@ This creates `integ-<your-name>-<stack-name>` and automatically:
 ```bash
 # Deploy base test stack
 ./scripts/local_deploy.sh base
-
 ```
 
 #### 3. Run Integration Tests
@@ -153,7 +85,19 @@ npm test
 
 # Run specific test file
 npm test -- base.test.ts
-
 ```
 
 **Note**: Tests wait several minutes after Lambda invocation to allow telemetry to propagate to Datadog.
+
+
+## Guidelines
+
+### Naming
+* An `identifier` is used to differentiate the different stacks. For local development, the identifier is automatically set using the command `whoami` and parses the user's first name. For gitlab pipelines, the identifier is the git commit short sha.
+* Stacks should be named `integ-<identifier>-<stack name>`
+* Lambda functions should be named `<stack-id>-<function name>`
+
+### Stacks
+* Stacks automatically get destroyed at the end of the gitlab integration tests step.
+* Stacks should be setup to not retain resources. A helper function `createLogGroup` exists with `removalPolicy: cdk.RemovalPolicy.DESTROY`. 
+* Stacks should include the tag `extension_integration_test: true`. This gets added in `app.ts`. This lets us easily run scripts to clean up old stacks in case the cleanup step was missed.

integration-tests/bin/app.ts

@@ -4,17 +4,23 @@ import * as cdk from 'aws-cdk-lib';
 import {Base} from '../lib/stacks/base';
 import {Otlp} from '../lib/stacks/otlp';
 import {Snapstart} from '../lib/stacks/snapstart';
-import {getIdentifier} from '../tests/utils/config';
+import {LambdaManagedInstancesStack} from '../lib/stacks/lmi';
+import {ACCOUNT, getIdentifier, REGION} from '../config';
+import {CapacityProviderStack} from "../lib/capacity-provider";
 
 const app = new cdk.App();
 
 const env = {
-    account: process.env.CDK_DEFAULT_ACCOUNT || process.env.AWS_ACCOUNT_ID,
-    region: process.env.CDK_DEFAULT_REGION || process.env.AWS_REGION || 'us-east-1',
+    account: ACCOUNT,
+    region: REGION,
 };
 
 const identifier = getIdentifier();
 
+// Use the same Lambda Managed Instance Capacity Provider for all LMI functions.
+// It is slow to create/destroy the related resources.
+new CapacityProviderStack(app, `integ-default-capacity-provider`, {env});
+
 const stacks = [
     new Base(app, `integ-${identifier}-base`, {
         env,
@@ -25,6 +31,9 @@ const stacks = [
     new Snapstart(app, `integ-${identifier}-snapstart`, {
         env,
     }),
+    new LambdaManagedInstancesStack(app, `integ-${identifier}-lmi`, {
+        env,
+    }),
 ]
 
 // Tag all stacks so we can easily clean them up

integration-tests/tests/utils/config.ts → integration-tests/config.ts

@@ -1,5 +1,9 @@
 import * as os from 'os';
 
+export const ACCOUNT = process.env.CDK_DEFAULT_ACCOUNT || process.env.AWS_ACCOUNT_ID;
+export const REGION = process.env.CDK_DEFAULT_REGION || process.env.AWS_REGION || 'us-east-1';
+
+
 export function getIdentifier(): string {
   if (process.env.IDENTIFIER) {
     return process.env.IDENTIFIER;

integration-tests/lambda/lmi-dotnet/Function.cs

@@ -0,0 +1,20 @@
+using Amazon.Lambda.Core;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+[assembly: LambdaSerializer(typeof(Amazon.Lambda.Serialization.SystemTextJson.DefaultLambdaJsonSerializer))]
+
+namespace Function
+{
+    public class Handler
+    {
+        public async Task<Dictionary<string, object>> FunctionHandler(Dictionary<string, object> input, ILambdaContext context)
+        {
+            context.Logger.LogLine("Hello World!");
+            return new Dictionary<string, object>
+            {
+                { "statusCode", 200 }
+            };
+        }
+    }
+}

integration-tests/lambda/lmi-dotnet/Function.csproj

@@ -0,0 +1,14 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFramework>net8.0</TargetFramework>
+    <GenerateRuntimeConfigurationFiles>true</GenerateRuntimeConfigurationFiles>
+    <AWSProjectType>Lambda</AWSProjectType>
+    <CopyLocalLockFileAssemblies>true</CopyLocalLockFileAssemblies>
+    <PublishReadyToRun>true</PublishReadyToRun>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Amazon.Lambda.Core" Version="2.2.0" />
+    <PackageReference Include="Amazon.Lambda.Serialization.SystemTextJson" Version="2.4.0" />
+  </ItemGroup>
+</Project>

integration-tests/lambda/lmi-java/pom.xml

@@ -0,0 +1,49 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>example</groupId>
+    <artifactId>managed-instances-java-lambda</artifactId>
+    <version>1.0.0</version>
+    <packaging>jar</packaging>
+
+    <name>Managed Instances Java Lambda</name>
+
+    <properties>
+        <maven.compiler.source>21</maven.compiler.source>
+        <maven.compiler.target>21</maven.compiler.target>
+        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-lambda-java-core</artifactId>
+            <version>1.2.3</version>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <version>3.5.0</version>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>shade</goal>
+                        </goals>
+                        <configuration>
+                            <finalName>function</finalName>
+                            <createDependencyReducedPom>false</createDependencyReducedPom>
+                        </configuration>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+</project>

integration-tests/lambda/lmi-java/src/main/java/example/Handler.java

@@ -0,0 +1,13 @@
+package example;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import java.util.Map;
+
+public class Handler implements RequestHandler<Map<String, Object>, Map<String, Object>> {
+    @Override
+    public Map<String, Object> handleRequest(Map<String, Object> event, Context context) {
+        context.getLogger().log("Hello World!");
+        return Map.of("statusCode", 200);
+    }
+}

integration-tests/lambda/lmi-node/index.js

@@ -0,0 +1,10 @@
+exports.handler = async (event, context) => {
+  console.log('Hello World!');
+  return {
+    statusCode: 200,
+    body: JSON.stringify({
+      message: 'Success from Lambda Managed Instance',
+      requestId: context.requestId
+    })
+  };
+};

integration-tests/lambda/lmi-python/lambda_function.py

@@ -0,0 +1,14 @@
+import logging
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+def handler(event, context):
+    logger.info('Hello World!')
+    return {
+        'statusCode': 200,
+        'body': {
+            'message': 'Success from Lambda Managed Instance',
+            'requestId': context.aws_request_id
+        }
+    }

integration-tests/lib/capacity-provider.ts

@@ -0,0 +1,80 @@
+import * as cdk from 'aws-cdk-lib';
+import { Construct } from 'constructs';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
+import * as ec2 from 'aws-cdk-lib/aws-ec2';
+import * as iam from 'aws-cdk-lib/aws-iam';
+
+/**
+ * AWS Lambda Managed Instances Stack
+ *
+ * This stack demonstrates AWS Lambda Managed Instances with:
+ * - VPC with NAT Gateway for Datadog Extension connectivity
+ * - Capacity Provider with ARM64 architecture
+ */
+export class CapacityProviderStack extends cdk.Stack {
+    constructor(scope: Construct, id: string, props?: cdk.StackProps) {
+        super(scope, id, props);
+
+        const vpc = new ec2.Vpc(this, `${id}-vpc`, {
+            maxAzs: 3,
+            natGateways: 1,
+            subnetConfiguration: [
+                {
+                    name: 'Private',
+                    subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS,
+                    cidrMask: 24,
+                },
+                {
+                    name: 'Public',
+                    subnetType: ec2.SubnetType.PUBLIC,
+                    cidrMask: 24,
+                }
+            ],
+            enableDnsHostnames: true,
+            enableDnsSupport: true
+        });
+
+        const securityGroup = new ec2.SecurityGroup(this, `${id}-security-group`, {
+            vpc: vpc,
+            description: 'Security group for Lambda Managed Instances',
+            allowAllOutbound: false
+        });
+
+        // Allow HTTPS outbound for Datadog Extension and AWS services
+        securityGroup.addEgressRule(
+            ec2.Peer.anyIpv4(),
+            ec2.Port.tcp(443),
+            'Allow HTTPS to Datadog (*.datadoghq.com) and AWS services'
+        );
+
+        const operatorRole = new iam.Role(this, `${id}-operator-role`, {
+            assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
+            managedPolicies: [
+                iam.ManagedPolicy.fromAwsManagedPolicyName('AWSLambdaManagedEC2ResourceOperator')
+            ],
+            description: 'Role for Lambda to manage EC2 instances in capacity provider'
+        });
+
+        const capacityProvider = new lambda.CapacityProvider(this, `${id}-cp`, {
+            capacityProviderName: `${id}-cp`,
+            subnets: vpc.selectSubnets({
+                subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS
+            }).subnets,
+            securityGroups: [securityGroup],
+            architectures: [lambda.Architecture.ARM_64],
+            maxVCpuCount: 80,
+            scalingOptions: lambda.ScalingOptions.auto(),
+            operatorRole: operatorRole
+        });
+
+        new cdk.CfnOutput(this, 'VpcId', {
+            value: vpc.vpcId,
+            description: 'VPC ID'
+        });
+        new cdk.CfnOutput(this, 'CapacityProviderArn', {
+            value: capacityProvider.capacityProviderArn,
+            description: 'ARN of the Capacity Provider'
+        });
+
+    }
+}