From f48f3e7111e8fb47c01adf729ceca87d7cd08b2c Mon Sep 17 00:00:00 2001
From: Luciano Mammino <lucianomammino@gmail.com>
Date: Fri, 25 Aug 2023 11:38:02 +0100
Subject: [PATCH] fix(typescript): ssm context is enriched correctly

---
 packages/core/index.d.ts     |  3 ++-
 packages/ssm/index.d.ts      |  4 +++-
 packages/ssm/index.test-d.ts | 23 +++++++++++++++++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/packages/core/index.d.ts b/packages/core/index.d.ts
index 69f71fd26..06cd14fad 100644
--- a/packages/core/index.d.ts
+++ b/packages/core/index.d.ts
@@ -50,7 +50,8 @@ export interface MiddlewareObj<
   TEvent = unknown,
   TResult = any,
   TErr = Error,
-  TContext extends LambdaContext = LambdaContext
+  TContext extends LambdaContext = LambdaContext,
+  TInternal = {}
 > {
   before?: MiddlewareFn<TEvent, TResult, TErr, TContext>
   after?: MiddlewareFn<TEvent, TResult, TErr, TContext>
diff --git a/packages/ssm/index.d.ts b/packages/ssm/index.d.ts
index 8ffd28703..19b70a7e6 100644
--- a/packages/ssm/index.d.ts
+++ b/packages/ssm/index.d.ts
@@ -10,6 +10,8 @@ interface Options<AwsSSMClient = SSMClient>
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 type ExtractSingles<T> = T extends `/${infer _}` ? never : T
 
+export type Internal<TOptions extends Options> = Record<ExtractSingles<keyof TOptions['fetchData']>, JsonValue>
+
 export type Context<TOptions extends Options | undefined> = TOptions extends {
   setToContext: true
 }
@@ -23,6 +25,6 @@ export type Context<TOptions extends Options | undefined> = TOptions extends {
 
 declare function ssm<TOptions extends Options> (
   options?: TOptions
-): middy.MiddlewareObj<unknown, any, Error, Context<TOptions>>
+): middy.MiddlewareObj<unknown, any, Error, Context<TOptions>, Internal<TOptions>>
 
 export default ssm
diff --git a/packages/ssm/index.test-d.ts b/packages/ssm/index.test-d.ts
index 14a53df78..6bbdb5053 100644
--- a/packages/ssm/index.test-d.ts
+++ b/packages/ssm/index.test-d.ts
@@ -3,6 +3,7 @@ import { SSMClient } from '@aws-sdk/client-ssm'
 import { captureAWSv3Client } from 'aws-xray-sdk'
 import { expectType } from 'tsd'
 import ssm, { Context } from '.'
+import { JsonValue } from 'type-fest'
 
 // use with default options
 expectType<middy.MiddlewareObj<unknown, any, Error, Context<undefined>>>(ssm())
@@ -24,3 +25,25 @@ const options = {
 expectType<middy.MiddlewareObj<unknown, any, Error, Context<typeof options>>>(
   ssm(options)
 )
+
+
+// checks that values fetched are actually enriching the context correctly (#1084)
+middy()
+  .use(
+    ssm({
+      fetchData: {
+        accessToken: '/dev/service_name/access_token', // single value
+        dbParams: '/dev/service_name/database/', // object of values, key for each path
+        defaults: '/dev/defaults'
+      },
+      setToContext: true
+    })
+  )
+  .before((request) => {
+    // checks that the context is correctly enriched in before
+    expectType<Record<"accessToken" | "dbParams" | "defaults", JsonValue>>(request.context)
+  })
+  .handler(async (req, context) => {
+    // checks that the context is correctly enriched in handler
+    expectType<Record<"accessToken" | "dbParams" | "defaults", JsonValue>>(context)
+  })