From a50628b9f5e74de1aac159240d4ac031ceaef91a Mon Sep 17 00:00:00 2001
From: "Kamat, Trivikram" <16024985+trivikr@users.noreply.github.com>
Date: Fri, 9 Aug 2024 20:50:49 +0000
Subject: [PATCH] test(endpoints): read serviceId from service model

---
 tests/endpoints-2.0/endpoints-integration.spec.ts | 13 ++++++-------
 tests/endpoints-2.0/integration-test-types.ts     |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/tests/endpoints-2.0/endpoints-integration.spec.ts b/tests/endpoints-2.0/endpoints-integration.spec.ts
index 86c672e8dfda3..caac767ec0495 100644
--- a/tests/endpoints-2.0/endpoints-integration.spec.ts
+++ b/tests/endpoints-2.0/endpoints-integration.spec.ts
@@ -3,7 +3,7 @@ import { EndpointParameters, EndpointV2 } from "@smithy/types";
 import * as fs from "fs";
 import * as path from "path";
 
-import { EndpointExpectation, EndpointTestCase, ServiceModel, ServiceNamespace } from "./integration-test-types";
+import { EndpointExpectation, ServiceModel, ServiceNamespace } from "./integration-test-types";
 
 describe("client list", () => {
   const root = path.join(__dirname, "..", "..");
@@ -40,7 +40,7 @@ describe("client list", () => {
       for (const value of Object.values(model.shapes)) {
         if (typeof value === "object" && value !== null && "type" in value && value.type === "service") {
           const service = value as ServiceModel;
-          runTestCases(service, namespace, defaultEndpointResolver, "");
+          runTestCases(service, namespace, defaultEndpointResolver);
           break;
         }
       }
@@ -53,12 +53,11 @@ describe("client list", () => {
 function runTestCases(
   service: ServiceModel,
   namespace: ServiceNamespace,
-  defaultEndpointResolver: (endpointParams: EndpointParameters) => EndpointV2,
-  serviceId: string
+  defaultEndpointResolver: (endpointParams: EndpointParameters) => EndpointV2
 ) {
-  const [, tests] = Object.entries(service.traits).find(([k, v]) => k === "smithy.rules#endpointTests") as any;
-  if (tests?.testCases) {
-    const testCases = tests.testCases as EndpointTestCase[];
+  const serviceId = service.traits["aws.api#service"].serviceId;
+  const testCases = service.traits["smithy.rules#endpointTests"]?.testCases;
+  if (testCases) {
     for (const testCase of testCases) {
       const { documentation, params = {}, expect: expectation, operationInputs } = testCase;
 
diff --git a/tests/endpoints-2.0/integration-test-types.ts b/tests/endpoints-2.0/integration-test-types.ts
index de9b29d6b06e5..1aad9e7109a10 100644
--- a/tests/endpoints-2.0/integration-test-types.ts
+++ b/tests/endpoints-2.0/integration-test-types.ts
@@ -38,7 +38,7 @@ export interface ServiceModel {
     "aws.api#service": {
       serviceId: string;
     };
-    "smithy.rules#endpointTests": {
+    "smithy.rules#endpointTests"?: {
       testCases: EndpointTestCase[];
     };
   };