fix: support Lambda instrumentation for contextManager: 'patch'; refactor Lambda tests (#3305)

Changes in #3285 broke Lambda instrumentation if the deprecated
`contextManager:'patch' config var was used. This restores support for that, FWIW.

Also, refactoring so we can drop AgentMock and TransactionMock under 
test/lambda/... to simplify:
- roll promises.test.js into the main test file
- roll trace-context.test.js into the main lambda test file
- roll transaction.test.js into main lambda test file
- rename main lambda test file to lambda.test.js
- drop now unused MockAgent

Author: trentm

docs/lambda.asciidoc

@@ -61,14 +61,14 @@ Your Lambda function invocations should be traced from now on.
 === Features
 
 The AWS Lambda instrumentation will report a transaction for all function invocations
-and trace any <<compatibility-frameworks,support modules>>. In addition, the
+and trace any <<compatibility-frameworks,supported modules>>. In addition, the
 created transactions will capture additional data for a number of Lambda
-trigger types -- API Gateway, SNS, SQS, S3 (with the trigger is a single event),
+trigger types -- API Gateway, SNS, SQS, S3 (when the trigger is a single event),
 and ELB.
 
 A transaction will be reported for Lambda invocations that fail due to a
 timeout, crash, `uncaughtException`, or `unhandledRejection`. (This requires
-APM agent v3.46.0 or later and
+APM agent v3.45.0 or later and
 https://www.elastic.co/guide/en/apm/lambda/current/aws-lambda-arch.html[Elastic's APM Lambda extension]
 version 1.4.0 or later.)
 
@@ -80,4 +80,5 @@ version 1.4.0 or later.)
 * System and custom metrics are not collected for Lambda functions. This is both because most of those are irrelevant
 and because the interval-based event sending model is not suitable for FaaS environments.
 * Lambda instrumentation does not currently work when the (deprecated) <<context-manager,`contextManager: 'patch'`>> configuration setting is used.
-* The APM agent does not yet support a Lambda handler module that uses ECMAScript modules (ESM). That means a your handler file name should end with ".js" (and not have `"type": "module"` in package.json if you have one) or end with ".cjs". A handler file that uses the ".mjs" suffix will not be instrumented by the APM agent.
+* The APM agent does not yet support a Lambda handler module that uses ECMAScript modules (ESM). That means your handler file name should end with ".js" (and not have `"type": "module"` in package.json if you have one) or end with ".cjs". A handler file that uses the ".mjs" suffix will not be instrumented by the APM agent.
+

lib/lambda.js

@@ -438,6 +438,7 @@ function setS3SingleData (trans, event, context, faasId, isColdStart) {
 
 function elasticApmAwsLambda (agent) {
   const log = agent.logger
+  const ins = agent._instrumentation
 
   /**
    * Register this transaction with the Lambda extension, if possible.  This
@@ -527,34 +528,43 @@ function elasticApmAwsLambda (agent) {
     })
   }
 
-  function wrapContext (trans, event, context, triggerType) {
+  function wrapContext (runContext, trans, event, context, triggerType) {
     shimmer.wrap(context, 'succeed', (origSucceed) => {
-      return function wrappedSucceed (result) {
-        endAndFlushTransaction(null, result, trans, event, context, triggerType, function () {
-          origSucceed(result)
-        })
-      }
+      return ins.bindFunctionToRunContext(
+        runContext,
+        function wrappedSucceed (result) {
+          endAndFlushTransaction(null, result, trans, event, context, triggerType, function () {
+            origSucceed(result)
+          })
+        }
+      )
     })
 
     shimmer.wrap(context, 'fail', (origFail) => {
-      return function wrappedFail (err) {
-        endAndFlushTransaction(err, null, trans, event, context, triggerType, function () {
-          origFail(err)
-        })
-      }
+      return ins.bindFunctionToRunContext(
+        runContext,
+        function wrappedFail (err) {
+          endAndFlushTransaction(err, null, trans, event, context, triggerType, function () {
+            origFail(err)
+          })
+        }
+      )
     })
 
     shimmer.wrap(context, 'done', (origDone) => {
-      return wrapLambdaCallback(trans, event, context, triggerType, origDone)
+      return wrapLambdaCallback(runContext, trans, event, context, triggerType, origDone)
     })
   }
 
-  function wrapLambdaCallback (trans, event, context, triggerType, callback) {
-    return function wrappedLambdaCallback (err, result) {
-      endAndFlushTransaction(err, result, trans, event, context, triggerType, () => {
-        callback(err, result)
-      })
-    }
+  function wrapLambdaCallback (runContext, trans, event, context, triggerType, callback) {
+    return ins.bindFunctionToRunContext(
+      runContext,
+      function wrappedLambdaCallback (err, result) {
+        endAndFlushTransaction(err, result, trans, event, context, triggerType, () => {
+          callback(err, result)
+        })
+      }
+    )
   }
 
   return function wrapLambdaHandler (type, fn) {
@@ -644,13 +654,14 @@ function elasticApmAwsLambda (agent) {
       // Note: Wrapping context needs to happen *before any `await` calls* in
       // this function, otherwise the Lambda Node.js Runtime will call the
       // *unwrapped* `context.{succeed,fail,done}()` methods.
-      wrapContext(trans, event, context, triggerType)
-      const wrappedCallback = wrapLambdaCallback(trans, event, context, triggerType, callback)
+      const transRunContext = ins.currRunContext()
+      wrapContext(transRunContext, trans, event, context, triggerType)
+      const wrappedCallback = wrapLambdaCallback(transRunContext, trans, event, context, triggerType, callback)
 
       await registerTransaction(trans, context.awsRequestId)
 
       try {
-        var retval = fn.call(this, event, context, wrappedCallback)
+        const retval = ins.withRunContext(transRunContext, fn, this, event, context, wrappedCallback)
         if (retval instanceof Promise) {
           return retval
         } else {