From e7367f87f407ff2d0cb163ed405d7783e24d9fd8 Mon Sep 17 00:00:00 2001
From: Yao Zhao <yaozhaoy@gmail.com>
Date: Tue, 8 Dec 2020 14:36:14 -0800
Subject: [PATCH] Fixed an issue that the resolved environment was not cached

---
 build.gradle                                  |  2 +-
 examples/lambda/build.gradle                  |  4 +++
 examples/lambda/src/main/resources/log4j2.xml | 31 +++++++++++++++++++
 .../emf/environment/EnvironmentProvider.java  | 10 +++---
 .../environment/EnvironmentProviderTest.java  | 24 +++++++++++++-
 5 files changed, 63 insertions(+), 8 deletions(-)
 create mode 100644 examples/lambda/src/main/resources/log4j2.xml

diff --git a/build.gradle b/build.gradle
index 827d85bd..9d7d0ee3 100644
--- a/build.gradle
+++ b/build.gradle
@@ -15,7 +15,7 @@
 
 plugins {
 	id 'java-library'
-	id 'com.diffplug.spotless' version '5.1.0'
+	id 'com.diffplug.spotless' version '5.8.2'
 	id 'maven-publish'
 	id 'signing'
 }
diff --git a/examples/lambda/build.gradle b/examples/lambda/build.gradle
index b5020878..0a217d1b 100644
--- a/examples/lambda/build.gradle
+++ b/examples/lambda/build.gradle
@@ -9,6 +9,10 @@ repositories {
 
 dependencies {
     implementation 'com.amazonaws:aws-lambda-java-core:1.2.1'
+    implementation  "org.apache.logging.log4j:log4j-api:2.13.3"
+    implementation "org.apache.logging.log4j:log4j-core:2.13.3"
+    implementation "org.apache.logging.log4j:log4j-slf4j-impl:2.13.3"
+    implementation 'com.amazonaws:aws-lambda-java-log4j2:1.2.0'
     implementation rootProject
         
 }
diff --git a/examples/lambda/src/main/resources/log4j2.xml b/examples/lambda/src/main/resources/log4j2.xml
new file mode 100644
index 00000000..2330ecba
--- /dev/null
+++ b/examples/lambda/src/main/resources/log4j2.xml
@@ -0,0 +1,31 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+  ~
+  ~   Licensed under the Apache License, Version 2.0 (the "License").
+  ~   You may not use this file except in compliance with the License.
+  ~   You may obtain a copy of the License at
+  ~
+  ~       http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~   Unless required by applicable law or agreed to in writing, software
+  ~   distributed under the License is distributed on an "AS IS" BASIS,
+  ~   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~   See the License for the specific language governing permissions and
+  ~   limitations under the License.
+  -->
+
+<Configuration packages="com.amazonaws.services.lambda.runtime.log4j2.LambdaAppender">
+    <Appenders>
+        <Lambda name="Lambda">
+            <PatternLayout>
+                <pattern>%d{yyyy-MM-dd HH:mm:ss} %X{AWSRequestId} %-5p %c{1}:%L - %m%n</pattern>
+            </PatternLayout>
+        </Lambda>
+    </Appenders>
+    <Loggers>
+        <Root level="debug">
+            <AppenderRef ref="Lambda" />
+        </Root>
+    </Loggers>
+</Configuration>
diff --git a/src/main/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProvider.java b/src/main/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProvider.java
index 7df3e9dc..49522b7b 100644
--- a/src/main/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProvider.java
+++ b/src/main/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProvider.java
@@ -54,12 +54,10 @@ public CompletableFuture<Environment> resolveEnvironment() {
         CompletableFuture<Optional<Environment>> resolvedEnv = discoverEnvironmentAsync();
 
         return resolvedEnv.thenApply(
-                optionalEnv ->
-                        optionalEnv.orElseGet(
-                                () -> {
-                                    cachedEnvironment = defaultEnvironment;
-                                    return cachedEnvironment;
-                                }));
+                optionalEnv -> {
+                    cachedEnvironment = optionalEnv.orElse(defaultEnvironment);
+                    return cachedEnvironment;
+                });
     }
 
     public Environment getDefaultEnvironment() {
diff --git a/src/test/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProviderTest.java b/src/test/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProviderTest.java
index 3994574e..3407ce15 100644
--- a/src/test/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProviderTest.java
+++ b/src/test/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProviderTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.*;
 import static org.powermock.api.mockito.PowerMockito.when;
 
 import com.github.javafaker.Faker;
@@ -82,6 +82,28 @@ public void testResolveEnvironmentReturnsLambdaEnvironment() {
         assertTrue(resolvedEnvironment.join() instanceof LambdaEnvironment);
     }
 
+    @Test
+    public void testResolveEnvironmentFromCacheAfterSecondCall() {
+        environmentProvider.cleanResolvedEnvironment();
+        String lambdaFunctionName = faker.name().name();
+
+        PowerMockito.mockStatic(SystemWrapper.class);
+        when(SystemWrapper.getenv("AWS_LAMBDA_FUNCTION_NAME")).thenReturn(lambdaFunctionName);
+
+        // First call
+        CompletableFuture<Environment> resolvedEnvironment =
+                environmentProvider.resolveEnvironment();
+
+        assertTrue(resolvedEnvironment.join() instanceof LambdaEnvironment);
+
+        // Second call
+        assertTrue(environmentProvider.resolveEnvironment().join() instanceof LambdaEnvironment);
+
+        // The SystemWrapper.getenv should only be called once due to cache
+        PowerMockito.verifyStatic(SystemWrapper.class);
+        SystemWrapper.getenv("AWS_LAMBDA_FUNCTION_NAME");
+    }
+
     @Test
     public void testResolveEnvironmentReturnsLambdaFromOverride() {
         PowerMockito.mockStatic(EnvironmentConfigurationProvider.class);