Clarify Tests, ensure it works with Bean injection

Author: mattcorey

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyServletContext.java

@@ -410,6 +410,9 @@ public String getVirtualServerName() {
         return null;
     }
 
+    public String getStage() {
+        return awsProxyRequest.getRequestContext().getStage();
+    }
 
     public static ServletContext getInstance(AwsProxyRequest request, Context lambdaContext) {
         if (instance == null) {
@@ -418,4 +421,8 @@ public static ServletContext getInstance(AwsProxyRequest request, Context lambda
 
         return instance;
     }
+
+    public static void clearServletContextCache() {
+        instance = null;
+    }
 }

aws-serverless-java-container-spring/src/main/java/com/amazonaws/serverless/proxy/spring/LambdaSpringApplicationInitializer.java

@@ -47,6 +47,7 @@ public class LambdaSpringApplicationInitializer implements WebApplicationInitial
     private ConfigurableWebApplicationContext applicationContext;
     private boolean refreshContext = true;
     private List<ServletContextListener> contextListeners;
+    private String lambdaStage;
 
     // Dynamically instantiated properties
     private ServletConfig dispatcherConfig;
@@ -84,8 +85,13 @@ public void dispatch(HttpServletRequest request, HttpServletResponse response)
         dispatcherServlet.service(request, response);
     }
 
+    public void setLambdaStage(String lambdaStage) {
+        this.lambdaStage = lambdaStage;
+    }
+
     @Override
     public void onStartup(ServletContext servletContext) throws ServletException {
+        applicationContext.getEnvironment().setActiveProfiles(lambdaStage);
         applicationContext.setServletContext(servletContext);
 
         dispatcherConfig = new DefaultDispatcherConfig(servletContext);

aws-serverless-java-container-spring/src/main/java/com/amazonaws/serverless/proxy/spring/SpringLambdaContainerHandler.java

@@ -16,10 +16,7 @@
 import com.amazonaws.serverless.proxy.internal.*;
 import com.amazonaws.serverless.proxy.internal.model.AwsProxyRequest;
 import com.amazonaws.serverless.proxy.internal.model.AwsProxyResponse;
-import com.amazonaws.serverless.proxy.internal.servlet.AwsHttpServletResponse;
-import com.amazonaws.serverless.proxy.internal.servlet.AwsProxyHttpServletRequest;
-import com.amazonaws.serverless.proxy.internal.servlet.AwsProxyHttpServletRequestReader;
-import com.amazonaws.serverless.proxy.internal.servlet.AwsProxyHttpServletResponseWriter;
+import com.amazonaws.serverless.proxy.internal.servlet.*;
 import com.amazonaws.services.lambda.runtime.Context;
 import org.springframework.web.context.ConfigurableWebApplicationContext;
 import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
@@ -112,6 +109,7 @@ protected void handleRequest(AwsProxyHttpServletRequest containerRequest, AwsHtt
 
         // wire up the application context on the first invocation
         if (!initialized) {
+            initializer.setLambdaStage(((AwsProxyServletContext) containerRequest.getServletContext()).getStage());
             ServletContext context = containerRequest.getServletContext();
             initializer.onStartup(context);
             initialized = true;

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/echoapp/EchoSpringAppConfig.java

@@ -2,9 +2,7 @@
 
 import com.amazonaws.serverless.exceptions.ContainerInitializationException;
 import com.amazonaws.serverless.proxy.internal.testutils.MockLambdaContext;
-import com.amazonaws.serverless.proxy.spring.LambdaSpringApplicationInitializer;
 import com.amazonaws.serverless.proxy.spring.SpringLambdaContainerHandler;
-import com.amazonaws.services.lambda.runtime.Context;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/echoapp/profile/DefaultProfileConfiguration.java

@@ -0,0 +1,12 @@
+package com.amazonaws.serverless.proxy.spring.echoapp.profile;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration
+public class DefaultProfileConfiguration {
+    @Bean
+    public String beanInjectedValue() {
+        return "default-profile-from-bean";
+    }
+}

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/echoapp/profile/DefaultProfileResource.java

@@ -1,6 +1,7 @@
 package com.amazonaws.serverless.proxy.spring.echoapp.profile;
 
 import com.amazonaws.serverless.proxy.spring.echoapp.model.MapResponseModel;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestMethod;
@@ -15,12 +16,16 @@ public class DefaultProfileResource {
     @Value("${spring-proxy.not-overridden-test}")
     private String noOverride;
 
+    @Autowired
+    private String beanInjectedValue;
+
     @RequestMapping(path = "/spring-properties", method = RequestMethod.GET)
     public MapResponseModel loadProperties() {
         MapResponseModel model = new MapResponseModel();
 
         model.addValue("profileTest", profileTest);
         model.addValue("noOverride", noOverride);
+        model.addValue("beanInjectedValue", beanInjectedValue);
         return model;
     }
 }

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/echoapp/profile/OverrideProfileConfiguration.java

@@ -0,0 +1,14 @@
+package com.amazonaws.serverless.proxy.spring.echoapp.profile;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Profile;
+
+@Configuration
+@Profile("override")
+public class OverrideProfileConfiguration {
+    @Bean
+    public String beanInjectedValue() {
+        return "override-profile-from-bean";
+    }
+}

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/profile/SpringProfileTest.java

@@ -2,61 +2,71 @@
 
 import com.amazonaws.serverless.proxy.internal.model.AwsProxyRequest;
 import com.amazonaws.serverless.proxy.internal.model.AwsProxyResponse;
+import com.amazonaws.serverless.proxy.internal.servlet.AwsProxyServletContext;
 import com.amazonaws.serverless.proxy.internal.testutils.AwsProxyRequestBuilder;
 import com.amazonaws.serverless.proxy.internal.testutils.MockLambdaContext;
 import com.amazonaws.serverless.proxy.spring.SpringLambdaContainerHandler;
 import com.amazonaws.serverless.proxy.spring.echoapp.EchoSpringAppConfig;
 import com.amazonaws.serverless.proxy.spring.echoapp.model.MapResponseModel;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.TestExecutionListeners;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.context.support.DependencyInjectionTestExecutionListener;
 import org.springframework.test.context.web.WebAppConfiguration;
 
-import java.io.IOException;
-
 import static org.junit.Assert.assertEquals;
 
-
 @RunWith(SpringJUnit4ClassRunner.class)
-@ContextConfiguration(classes = {EchoSpringAppConfig.class})
+@ContextConfiguration(classes = EchoSpringAppConfig.class)
 @WebAppConfiguration
+@TestExecutionListeners(inheritListeners = false, listeners = {DependencyInjectionTestExecutionListener.class})
 public class SpringProfileTest {
     @Autowired
     private ObjectMapper objectMapper;
 
     @Autowired
     private MockLambdaContext lambdaContext;
 
-    @Autowired
-    private SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> handler;
+    @Before
+    public void clearServletContextCache() {
+        AwsProxyServletContext.clearServletContextCache();
+    }
 
     @Test
-    public void profile_defaultProfile() throws IOException {
+    public void profile_defaultProfile() throws Exception {
         AwsProxyRequest request = new AwsProxyRequestBuilder("/profile/spring-properties", "GET")
                 .build();
 
+        SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> handler = SpringLambdaContainerHandler.getAwsProxyHandler(EchoSpringAppConfig.class);
         AwsProxyResponse output = handler.proxy(request, lambdaContext);
         assertEquals(200, output.getStatusCode());
 
         MapResponseModel response = objectMapper.readValue(output.getBody(), MapResponseModel.class);
+        assertEquals(3, response.getValues().size());
         assertEquals("default-profile", response.getValues().get("profileTest"));
         assertEquals("not-overridden", response.getValues().get("noOverride"));
+        assertEquals("default-profile-from-bean", response.getValues().get("beanInjectedValue"));
     }
 
     @Test
-    public void profile_overrideProfile() throws IOException {
+    public void profile_overrideProfile() throws Exception {
         AwsProxyRequest request = new AwsProxyRequestBuilder("/profile/spring-properties", "GET")
                 .stage("override")
                 .build();
 
+        SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> handler = SpringLambdaContainerHandler.getAwsProxyHandler(EchoSpringAppConfig.class);
         AwsProxyResponse output = handler.proxy(request, lambdaContext);
         assertEquals(200, output.getStatusCode());
 
         MapResponseModel response = objectMapper.readValue(output.getBody(), MapResponseModel.class);
+        assertEquals(3, response.getValues().size());
         assertEquals("override-profile", response.getValues().get("profileTest"));
         assertEquals("not-overridden", response.getValues().get("noOverride"));
+        assertEquals("override-profile-from-bean", response.getValues().get("beanInjectedValue"));
     }
 }