fix: call command.resolveMiddleware instead of client.send

refs: trivikr#105

Author: trivikr

clients/client-timestream-query/TimestreamQueryClient.ts

@@ -223,10 +223,8 @@ export class TimestreamQueryClient extends __Client<
     let _config_4 = resolveHostHeaderConfig(_config_3);
     let _config_5 = resolveAwsAuthConfig(_config_4);
     let _config_6 = resolveUserAgentConfig(_config_5);
-    let _config_7 = resolveEndpointDiscoveryConfig(_config_6);
+    let _config_7 = resolveEndpointDiscoveryConfig(_config_6, DescribeEndpointsCommand);
     super(_config_7);
-    _config_7.client = this;
-    _config_7.endpointDiscoveryCommandCtor = DescribeEndpointsCommand;
     this.config = _config_7;
     this.middlewareStack.use(getRetryPlugin(this.config));
     this.middlewareStack.use(getContentLengthPlugin(this.config));

clients/client-timestream-query/commands/QueryCommand.ts

@@ -55,10 +55,13 @@ export class QueryCommand extends $Command<QueryCommandInput, QueryCommandOutput
     configuration: TimestreamQueryClientResolvedConfig,
     options?: __HttpHandlerOptions
   ): Handler<QueryCommandInput, QueryCommandOutput> {
+    const isDiscoveredEndpointRequired = true;
     this.middlewareStack.use(getSerdePlugin(configuration, this.serialize, this.deserialize));
     this.middlewareStack.use(
       getEndpointDiscoveryCommandPlugin(configuration, {
-        isDiscoveredEndpointRequired: true,
+        isDiscoveredEndpointRequired,
+        clientStack,
+        options,
       })
     );
 

packages/middleware-endpoint-discovery/src/getEndpointDiscoveryCommandPlugin.ts

@@ -1,4 +1,4 @@
-import { FinalizeRequestHandlerOptions, Pluggable } from "@aws-sdk/types";
+import { FinalizeRequestHandlerOptions, HttpHandlerOptions, MiddlewareStack, Pluggable } from "@aws-sdk/types";
 
 import { endpointDiscoveryMiddleware } from "./endpointDiscoveryMiddleware";
 import { EndpointDiscoveryResolvedConfig } from "./resolveEndpointDiscoveryConfig";
@@ -12,6 +12,8 @@ export const endpointDiscoveryMiddlewareOptions: FinalizeRequestHandlerOptions =
 
 export type EndpointDiscoveryMiddlewareConfig = {
   isDiscoveredEndpointRequired: boolean;
+  clientStack: MiddlewareStack<any, any>;
+  options?: HttpHandlerOptions;
   identifiers?: { [key: string]: string };
 };
 

packages/middleware-endpoint-discovery/src/resolveEndpointDiscoveryConfig.spec.ts

@@ -5,56 +5,58 @@ import { resolveEndpointDiscoveryConfig } from "./resolveEndpointDiscoveryConfig
 jest.mock("@aws-sdk/endpoint-cache");
 
 describe(resolveEndpointDiscoveryConfig.name, () => {
-  const isCustomEndpoint = false;
-  const credentials = jest.fn();
-  const endpointDiscoveryEnabledProvider = jest.fn().mockResolvedValue(undefined);
+  const endpointDiscoveryCommandCtor = jest.fn();
+  const mockInput = {
+    isCustomEndpoint: false,
+    credentials: jest.fn(),
+    endpointDiscoveryEnabledProvider: jest.fn(),
+  };
 
   afterEach(() => {
     jest.clearAllMocks();
   });
 
+  it("assigns endpointDiscoveryCommandCtor in resolvedConfig", () => {
+    const resolvedConfig = resolveEndpointDiscoveryConfig(mockInput, endpointDiscoveryCommandCtor);
+    expect(resolvedConfig.endpointDiscoveryCommandCtor).toStrictEqual(endpointDiscoveryCommandCtor);
+  });
+
   describe("endpointCache", () => {
     it("creates cache of size endpointCacheSize if passed", () => {
       const endpointCacheSize = 100;
-      resolveEndpointDiscoveryConfig({
-        isCustomEndpoint,
-        credentials,
-        endpointCacheSize,
-        endpointDiscoveryEnabledProvider,
-      });
+      resolveEndpointDiscoveryConfig(
+        {
+          ...mockInput,
+          endpointCacheSize,
+        },
+        endpointDiscoveryCommandCtor
+      );
       expect(EndpointCache).toBeCalledWith(endpointCacheSize);
     });
 
     it("creates cache of size 1000 if endpointCacheSize not passed", () => {
-      resolveEndpointDiscoveryConfig({
-        isCustomEndpoint,
-        credentials,
-        endpointDiscoveryEnabledProvider,
-      });
+      resolveEndpointDiscoveryConfig(mockInput, endpointDiscoveryCommandCtor);
       expect(EndpointCache).toBeCalledWith(1000);
     });
   });
 
   describe("endpointDiscoveryEnabled", () => {
     it.each<boolean>([false, true])(`sets to value passed in the config: %s`, (endpointDiscoveryEnabled) => {
-      const resolvedConfig = resolveEndpointDiscoveryConfig({
-        isCustomEndpoint,
-        credentials,
-        endpointDiscoveryEnabled,
-        endpointDiscoveryEnabledProvider,
-      });
+      const resolvedConfig = resolveEndpointDiscoveryConfig(
+        {
+          ...mockInput,
+          endpointDiscoveryEnabled,
+        },
+        endpointDiscoveryCommandCtor
+      );
       expect(resolvedConfig.endpointDiscoveryEnabled()).resolves.toBe(endpointDiscoveryEnabled);
-      expect(endpointDiscoveryEnabledProvider).not.toHaveBeenCalled();
+      expect(mockInput.endpointDiscoveryEnabledProvider).not.toHaveBeenCalled();
       expect(resolvedConfig.isClientEndpointDiscoveryEnabled).toStrictEqual(true);
     });
 
     it(`sets to endpointDiscoveryEnabledProvider if value is not passed`, () => {
-      const resolvedConfig = resolveEndpointDiscoveryConfig({
-        isCustomEndpoint,
-        credentials,
-        endpointDiscoveryEnabledProvider,
-      });
-      expect(resolvedConfig.endpointDiscoveryEnabled).toBe(endpointDiscoveryEnabledProvider);
+      const resolvedConfig = resolveEndpointDiscoveryConfig(mockInput, endpointDiscoveryCommandCtor);
+      expect(resolvedConfig.endpointDiscoveryEnabled).toBe(mockInput.endpointDiscoveryEnabledProvider);
       expect(resolvedConfig.isClientEndpointDiscoveryEnabled).toStrictEqual(false);
     });
   });

packages/middleware-endpoint-discovery/src/resolveEndpointDiscoveryConfig.ts

@@ -25,17 +25,18 @@ interface PreviouslyResolved {
 export interface EndpointDiscoveryResolvedConfig {
   isCustomEndpoint: boolean;
   credentials: Provider<Credentials>;
-  client?: Client<any, any, any>;
-  endpointDiscoveryCommandCtor?: new (comandConfig: any) => Command<any, any, any, any, any>;
+  endpointDiscoveryCommandCtor: new (comandConfig: any) => Command<any, any, any, any, any>;
   endpointCache: EndpointCache;
   endpointDiscoveryEnabled: Provider<boolean | undefined>;
   isClientEndpointDiscoveryEnabled: boolean;
 }
 
 export const resolveEndpointDiscoveryConfig = <T>(
-  input: T & PreviouslyResolved & EndpointDiscoveryInputConfig
+  input: T & PreviouslyResolved & EndpointDiscoveryInputConfig,
+  endpointDiscoveryCommandCtor: new (comandConfig: any) => Command<any, any, any, any, any>
 ): T & EndpointDiscoveryResolvedConfig => ({
   ...input,
+  endpointDiscoveryCommandCtor,
   endpointCache: new EndpointCache(input.endpointCacheSize ?? 1000),
   endpointDiscoveryEnabled:
     input.endpointDiscoveryEnabled !== undefined

packages/middleware-endpoint-discovery/src/updateDiscoveredEndpointInCache.spec.ts

@@ -8,22 +8,20 @@ describe(updateDiscoveredEndpointInCache.name, () => {
   const mockGet = jest.fn();
   const mockSet = jest.fn();
   const mockDelete = jest.fn();
-  const mockSend = jest.fn();
 
-  const mockEndpoints = [{ Address: "mockAddress", CachePeriodInMinutes: 1 }];
+  const mockHandler = jest.fn();
+  const mockResolveMiddleware = jest.fn().mockReturnValue(mockHandler);
+
+  const mockEndpoints = [{ Address: "mockAddress", CachePeriodInMinutes: 2 }];
   const placeholderEndpoints = [{ Address: "", CachePeriodInMinutes: 1 }];
 
   const config = {
-    client: {
-      send: mockSend,
-      config: {},
-    },
     endpointCache: { get: mockGet, set: mockSet, delete: mockDelete },
   };
 
   const options = {
     commandName: "ExampleCommand",
-    endpointDiscoveryCommandCtor: jest.fn(),
+    endpointDiscoveryCommandCtor: jest.fn().mockReturnValue({ resolveMiddleware: mockResolveMiddleware }),
     isDiscoveredEndpointRequired: false,
     identifiers: { key: "value" },
   };
@@ -84,11 +82,11 @@ describe(updateDiscoveredEndpointInCache.name, () => {
         Operation: options.commandName.substr(0, options.commandName.length - 7),
         Identifiers: options.identifiers,
       });
-      expect(mockSend).toHaveBeenCalledTimes(1);
+      expect(mockHandler).toHaveBeenCalledTimes(1);
     };
 
     it("on successful call: updates cache", async () => {
-      mockSend.mockResolvedValueOnce({ Endpoints: mockEndpoints });
+      mockHandler.mockResolvedValueOnce({ output: { Endpoints: mockEndpoints } });
 
       // @ts-ignore
       await updateDiscoveredEndpointInCache(config, options);
@@ -103,7 +101,7 @@ describe(updateDiscoveredEndpointInCache.name, () => {
     describe("on error", () => {
       it(`throws if isDiscoveredEndpointRequired=true`, async () => {
         const error = new Error("rejected");
-        mockSend.mockRejectedValueOnce(error);
+        mockHandler.mockRejectedValueOnce(error);
 
         try {
           // @ts-ignore
@@ -128,7 +126,7 @@ describe(updateDiscoveredEndpointInCache.name, () => {
 
       it(`sets placeholder enpoint if isDiscoveredEndpointRequired=false`, async () => {
         const error = new Error("rejected");
-        mockSend.mockRejectedValueOnce(error);
+        mockHandler.mockRejectedValueOnce(error);
 
         // @ts-ignore
         await updateDiscoveredEndpointInCache(config, options);
@@ -150,7 +148,7 @@ describe(updateDiscoveredEndpointInCache.name, () => {
 
         it(`InvalidEndpointException`, async () => {
           const error = Object.assign(new Error("Invalid endpoint!"), { name: "InvalidEndpointException" });
-          mockSend.mockRejectedValueOnce(error);
+          mockHandler.mockRejectedValueOnce(error);
 
           // @ts-ignore
           await updateDiscoveredEndpointInCache(config, options);
@@ -160,7 +158,7 @@ describe(updateDiscoveredEndpointInCache.name, () => {
 
         it(`Status code: 421`, async () => {
           const error = Object.assign(new Error("Invalid endpoint!"), { $metadata: { httpStatusCode: 421 } });
-          mockSend.mockRejectedValueOnce(error);
+          mockHandler.mockRejectedValueOnce(error);
 
           // @ts-ignore
           await updateDiscoveredEndpointInCache(config, options);

packages/middleware-endpoint-discovery/src/updateDiscoveredEndpointInCache.ts

@@ -15,7 +15,7 @@ export const updateDiscoveredEndpointInCache = async (
   config: EndpointDiscoveryResolvedConfig,
   options: updateDiscoveredEndpointInCacheOptions
 ) => {
-  const { client, endpointCache } = config;
+  const { endpointCache } = config;
   const { commandName, identifiers } = options;
   const cacheKey = await getCacheKey(commandName, config, { identifiers });
 
@@ -41,8 +41,9 @@ export const updateDiscoveredEndpointInCache = async (
         Operation: commandName.substr(0, commandName.length - 7), // strip "Command"
         Identifiers: identifiers,
       });
-      const { Endpoints } = await client?.send(command);
-      endpointCache.set(cacheKey, Endpoints);
+      const handler = command.resolveMiddleware(options.clientStack, config, options.options);
+      const result = await handler(command);
+      endpointCache.set(cacheKey, result.output.Endpoints);
     } catch (error) {
       if (error.name === "InvalidEndpointException" || error.$metadata?.httpStatusCode === 421) {
         // Endpoint is invalid, delete the cache entry.