From 4d1f3b32474bfdaafc619427e8c796c314f5570c Mon Sep 17 00:00:00 2001
From: Armando Segnini <segnina@amazon.com>
Date: Tue, 10 Sep 2024 16:46:01 +0200
Subject: [PATCH] Fixed test kafkaApi e2e test

---
 framework/src/streaming/lib/msk/kafka-api.ts  |  2 ++
 .../src/streaming/lib/msk/msk-provisioned.ts  |  2 +-
 .../src/streaming/lib/msk/msk-serverless.ts   |  2 +-
 .../kafka-api-byo-service-token.e2e.test.ts   | 30 ++++++++++++-------
 framework/test/e2e/kafka-api.e2e.test.ts      |  6 +---
 5 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/framework/src/streaming/lib/msk/kafka-api.ts b/framework/src/streaming/lib/msk/kafka-api.ts
index 6b2f6e95..b641a710 100644
--- a/framework/src/streaming/lib/msk/kafka-api.ts
+++ b/framework/src/streaming/lib/msk/kafka-api.ts
@@ -124,6 +124,7 @@ export class KafkaApi extends TrackedConstruct {
         this.mskAclLogGroup = mskAclProvider.onEventHandlerLogGroup;
         this.mskAclFunction = mskAclProvider.onEventHandlerFunction;
         this.mskAclSecurityGroup = mskAclProvider.securityGroups;
+        this.serviceToken = mskAclProvider.serviceToken;
       }
     }
 
@@ -147,6 +148,7 @@ export class KafkaApi extends TrackedConstruct {
         this.mskIamLogGroup = mskIamProvider.onEventHandlerLogGroup;
         this.mskIamFunction = mskIamProvider.onEventHandlerFunction;
         this.mskIamSecurityGroup = mskIamProvider.securityGroups;
+        this.serviceToken = mskIamProvider.serviceToken;
       }
     }
 
diff --git a/framework/src/streaming/lib/msk/msk-provisioned.ts b/framework/src/streaming/lib/msk/msk-provisioned.ts
index 588e515c..a56d2d09 100644
--- a/framework/src/streaming/lib/msk/msk-provisioned.ts
+++ b/framework/src/streaming/lib/msk/msk-provisioned.ts
@@ -643,7 +643,7 @@ export class MskProvisioned extends TrackedConstruct {
     }
     
     new CfnOutput(this, 'ServiceToken', { 
-      value: this.kafkaApi.serviceToken!,
+      value: this.serviceToken!,
       exportName: 'ServiceToken',
     });
 
diff --git a/framework/src/streaming/lib/msk/msk-serverless.ts b/framework/src/streaming/lib/msk/msk-serverless.ts
index 52515a1b..c2bb493b 100644
--- a/framework/src/streaming/lib/msk/msk-serverless.ts
+++ b/framework/src/streaming/lib/msk/msk-serverless.ts
@@ -120,7 +120,7 @@ export class MskServerless extends TrackedConstruct {
     this.serviceToken = this.kafkaApi.serviceToken;
 
     new CfnOutput(this, 'ServiceToken', { 
-      value: this.kafkaApi.serviceToken!,
+      value: this.serviceToken!,
       exportName: 'ServiceToken',
     });
   }
diff --git a/framework/test/e2e/kafka-api-byo-service-token.e2e.test.ts b/framework/test/e2e/kafka-api-byo-service-token.e2e.test.ts
index a613beba..064e029a 100644
--- a/framework/test/e2e/kafka-api-byo-service-token.e2e.test.ts
+++ b/framework/test/e2e/kafka-api-byo-service-token.e2e.test.ts
@@ -10,12 +10,9 @@
 
 import { App, RemovalPolicy, CfnOutput } from 'aws-cdk-lib';
 
-import { CertificateAuthority } from 'aws-cdk-lib/aws-acmpca';
 import { SecurityGroup, SubnetType } from 'aws-cdk-lib/aws-ec2';
-import { CfnCluster } from 'aws-cdk-lib/aws-msk';
-import { Secret } from 'aws-cdk-lib/aws-secretsmanager';
 import { TestStack } from './test-stack';
-import { AclOperationTypes, AclPermissionTypes, AclResourceTypes, ClientAuthentication, KafkaApi, KafkaClientLogLevel, MskClusterType, MskServerless, ResourcePatternTypes } from '../../src/streaming';
+import { ClientAuthentication, KafkaApi, KafkaClientLogLevel, MskClusterType, MskServerless, Authentication} from '../../src/streaming';
 import { DataVpc, Utils } from '../../src/utils';
 
 
@@ -43,7 +40,7 @@ const msk = new MskServerless(stack, 'cluster', {
     kafkaClientLogLevel: KafkaClientLogLevel.DEBUG,
 });
 
-const mskApi = new KafkaApi(stack, 'kafkaApi', {
+const kafkaApi = new KafkaApi(stack, 'kafkaApi', {
     vpc: vpc.vpc,
     clusterArn: msk.cluster.attrArn,
     subnets: vpc.vpc.selectSubnets({ subnetType: SubnetType.PRIVATE_WITH_EGRESS }),
@@ -55,10 +52,16 @@ const mskApi = new KafkaApi(stack, 'kafkaApi', {
     serviceToken: msk.serviceToken,
 });
 
-msk.addTopic('topicServerelss', {
-    topic: 'dummy',
-    numPartitions: 1,
-}, RemovalPolicy.DESTROY, false, 1500);
+
+kafkaApi.setTopic('dummyTopic',
+    Authentication.IAM,
+    {
+      topic: 'dummy',
+      numPartitions: 3,
+    },
+    RemovalPolicy.DESTROY,
+    true, 1000
+);
 
 new CfnOutput(stack, 'clusterArn', {
     value: msk.cluster.attrArn,
@@ -71,11 +74,18 @@ beforeAll(async () => {
     deployResult = await testStack.deploy();
 }, 10000000);
 
-it('Containers runtime created successfully', async () => {
+test('MSK cluster created successfully', async () => {
     // THEN
     expect(deployResult.clusterArn).toContain('arn');
 });
 
+/*
+test('Kafka API outputs service token successfully', async () => {
+    // THEN
+    expect(deployResult.ServiceToken).toContain('arn');
+});
+*/
+
 afterAll(async () => {
     await testStack.destroy();
 }, 10000000);
\ No newline at end of file
diff --git a/framework/test/e2e/kafka-api.e2e.test.ts b/framework/test/e2e/kafka-api.e2e.test.ts
index 355fe795..7085b716 100644
--- a/framework/test/e2e/kafka-api.e2e.test.ts
+++ b/framework/test/e2e/kafka-api.e2e.test.ts
@@ -112,10 +112,6 @@ new CfnOutput(stack, 'clusterArn', {
   value: cfnCluster.attrArn,
 });
 
-new CfnOutput(stack, 'serviceToken', {
-  value: mskApi.serviceToken!,
-});
-
 let deployResult: Record<string, string>;
 
 beforeAll(async() => {
@@ -123,7 +119,7 @@ beforeAll(async() => {
   deployResult = await testStack.deploy();
 }, 10000000);
 
-it('Containers runtime created successfully', async () => {
+test('CMSK cluster created successfully', async () => {
   // THEN
   expect(deployResult.clusterArn).toContain('arn');
 });