From 3773c09e6275085987204e898e97235a12167ee2 Mon Sep 17 00:00:00 2001
From: Adam Ruka <adamruka@amazon.com>
Date: Mon, 3 Jan 2022 15:40:07 -0800
Subject: [PATCH] fix(lambda): imported Function still has region and account
 from its Stack, instead of its ARN

Fixes #18228
---
 .../test/web-distribution.test.ts             |  2 +-
 packages/@aws-cdk/aws-lambda/lib/function.ts  |  4 ++-
 .../@aws-cdk/aws-lambda/test/function.test.ts | 29 +++++++++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/packages/@aws-cdk/aws-cloudfront/test/web-distribution.test.ts b/packages/@aws-cdk/aws-cloudfront/test/web-distribution.test.ts
index d7612633cf93c..df1d401ae7de0 100644
--- a/packages/@aws-cdk/aws-cloudfront/test/web-distribution.test.ts
+++ b/packages/@aws-cdk/aws-cloudfront/test/web-distribution.test.ts
@@ -1377,7 +1377,7 @@ added the ellipsis so a user would know there was more to ...`,
     // GIVEN
     const stack = new cdk.Stack();
     const sourceBucket = new s3.Bucket(stack, 'Bucket');
-    const lambdaVersion = lambda.Version.fromVersionArn(stack, 'Version', 'arn:my-version');
+    const lambdaVersion = lambda.Version.fromVersionArn(stack, 'Version', 'arn:aws:lambda:function-region:111111111111:function:function-name');
 
     // WHEN
     new CloudFrontWebDistribution(stack, 'MyDistribution', {
diff --git a/packages/@aws-cdk/aws-lambda/lib/function.ts b/packages/@aws-cdk/aws-lambda/lib/function.ts
index 3b5df3c5c5c41..2b47aa8d7bca3 100644
--- a/packages/@aws-cdk/aws-lambda/lib/function.ts
+++ b/packages/@aws-cdk/aws-lambda/lib/function.ts
@@ -455,7 +455,9 @@ export class Function extends FunctionBase {
       protected readonly canCreatePermissions = attrs.sameEnvironment ?? this._isStackAccount();
 
       constructor(s: Construct, i: string) {
-        super(s, i);
+        super(s, i, {
+          environmentFromArn: functionArn,
+        });
 
         this.grantPrincipal = role || new iam.UnknownPrincipal({ resource: this });
 
diff --git a/packages/@aws-cdk/aws-lambda/test/function.test.ts b/packages/@aws-cdk/aws-lambda/test/function.test.ts
index 7b8902698de9b..04a8ab862b5b7 100644
--- a/packages/@aws-cdk/aws-lambda/test/function.test.ts
+++ b/packages/@aws-cdk/aws-lambda/test/function.test.ts
@@ -297,6 +297,35 @@ describe('function', () => {
     expect(imported.functionName).toEqual('ProcessKinesisRecords');
   });
 
+  describe('Function.fromFunctionAttributes()', () => {
+    let stack: cdk.Stack;
+
+    beforeEach(() => {
+      const app = new cdk.App();
+      stack = new cdk.Stack(app, 'Base', {
+        env: { account: '111111111111', region: 'stack-region' },
+      });
+    });
+
+    describe('for a function in a different account and region', () => {
+      let func: lambda.IFunction;
+
+      beforeEach(() => {
+        func = lambda.Function.fromFunctionAttributes(stack, 'iFunc', {
+          functionArn: 'arn:aws:lambda:function-region:222222222222:function:function-name',
+        });
+      });
+
+      test("the function's region is taken from the ARN", () => {
+        expect(func.env.region).toBe('function-region');
+      });
+
+      test("the function's account is taken from the ARN", () => {
+        expect(func.env.account).toBe('222222222222');
+      });
+    });
+  });
+
   describe('addPermissions', () => {
     test('imported Function w/ resolved account and function arn', () => {
       // GIVEN