Make origin aware system property source

Closes gh-9252

Author: mbhave

spring-boot/src/main/java/org/springframework/boot/env/OriginTrackedSystemPropertySource.java

@@ -16,12 +16,11 @@
 
 package org.springframework.boot.env;
 
-import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.Map;
 
 import org.springframework.boot.SpringApplication;
-import org.springframework.boot.origin.OriginTrackedValue;
+import org.springframework.boot.origin.Origin;
+import org.springframework.boot.origin.OriginLookup;
 import org.springframework.boot.origin.SystemEnvironmentOrigin;
 import org.springframework.core.Ordered;
 import org.springframework.core.env.ConfigurableEnvironment;
@@ -32,7 +31,7 @@
 /**
  * An {@link EnvironmentPostProcessor} that replaces the systemEnvironment
  * {@link SystemEnvironmentPropertySource} with an
- * {@link OriginTrackedSystemPropertySource} that can track the
+ * {@link OriginAwareSystemEnvironmentPropertySource} that can track the
  * {@link SystemEnvironmentOrigin} for every system environment property.
  *
  * @author Madhura Bhave
@@ -63,17 +62,11 @@ public void postProcessEnvironment(ConfigurableEnvironment environment,
 	@SuppressWarnings("unchecked")
 	private void replacePropertySource(ConfigurableEnvironment environment,
 			String sourceName, PropertySource<?> propertySource) {
-		if (propertySource.getSource() instanceof Map) {
 			Map<String, Object> originalSource = (Map<String, Object>) propertySource
 					.getSource();
-			Map<String, Object> originTrackedSource = new LinkedHashMap<>(originalSource);
-			originTrackedSource.entrySet().forEach(e -> e.setValue(OriginTrackedValue
-					.of(e.getValue(), new SystemEnvironmentOrigin(e.getKey()))));
-			OriginTrackedSystemPropertySource source = new OriginTrackedSystemPropertySource(
-					sourceName, Collections.unmodifiableMap(originTrackedSource));
+			SystemEnvironmentPropertySource source = new OriginAwareSystemEnvironmentPropertySource(sourceName, originalSource);
 			environment.getPropertySources().replace(sourceName, source);
 		}
-	}
 
 	@Override
 	public int getOrder() {
@@ -84,4 +77,24 @@ public void setOrder(int order) {
 		this.order = order;
 	}
 
+	/**
+	 * {@link SystemEnvironmentPropertySource} that also tracks {@link Origin}.
+	 */
+	protected static class OriginAwareSystemEnvironmentPropertySource extends SystemEnvironmentPropertySource
+			implements OriginLookup<String> {
+
+		OriginAwareSystemEnvironmentPropertySource(String name, Map<String, Object> source) {
+			super(name, source);
+		}
+
+		@Override
+		public Origin getOrigin(String key) {
+			String property = resolvePropertyName(key);
+			if (super.containsProperty(property)) {
+				return new SystemEnvironmentOrigin(property);
+			}
+			return null;
+		}
+	}
+
 }

spring-boot/src/main/java/org/springframework/boot/env/SystemEnvironmentPropertySourceEnvironmentPostProcessor.java

@@ -16,16 +16,18 @@
 
 package org.springframework.boot.env;
 
+import java.util.Collections;
 import java.util.Map;
 
 import org.junit.Before;
 import org.junit.Test;
 
-import org.springframework.boot.origin.OriginTrackedValue;
+import org.springframework.boot.env.SystemEnvironmentPropertySourceEnvironmentPostProcessor.OriginAwareSystemEnvironmentPropertySource;
 import org.springframework.boot.origin.SystemEnvironmentOrigin;
 import org.springframework.core.env.ConfigurableEnvironment;
 import org.springframework.core.env.PropertySource;
 import org.springframework.core.env.StandardEnvironment;
+import org.springframework.core.env.SystemEnvironmentPropertySource;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -50,27 +52,49 @@ public void postProcessShouldReplaceSystemEnvironmentPropertySource()
 		postProcessor.postProcessEnvironment(this.environment, null);
 		PropertySource<?> replaced = this.environment.getPropertySources()
 				.get("systemEnvironment");
-		assertThat(replaced).isInstanceOf(OriginTrackedSystemPropertySource.class);
+		assertThat(replaced).isInstanceOf(OriginAwareSystemEnvironmentPropertySource.class);
 	}
 
 	@Test
 	@SuppressWarnings("unchecked")
-	public void replacedPropertySourceShouldHaveOriginTrackedValues() throws Exception {
+	public void replacedPropertySourceShouldBeOriginAware() throws Exception {
 		SystemEnvironmentPropertySourceEnvironmentPostProcessor postProcessor = new SystemEnvironmentPropertySourceEnvironmentPostProcessor();
 		PropertySource<?> original = this.environment.getPropertySources()
 				.get("systemEnvironment");
 		postProcessor.postProcessEnvironment(this.environment, null);
-		PropertySource<?> replaced = this.environment.getPropertySources()
+		OriginAwareSystemEnvironmentPropertySource replaced = (OriginAwareSystemEnvironmentPropertySource) this.environment.getPropertySources()
 				.get("systemEnvironment");
 		Map<String, Object> originalMap = (Map<String, Object>) original.getSource();
-		Map<String, OriginTrackedValue> replacedMap = (Map<String, OriginTrackedValue>) replaced
+		Map<String, Object> replacedMap = replaced
 				.getSource();
 		for (Map.Entry<String, Object> entry : originalMap.entrySet()) {
-			OriginTrackedValue actual = replacedMap.get(entry.getKey());
-			assertThat(actual.getValue()).isEqualTo(entry.getValue());
-			assertThat(actual.getOrigin())
-					.isEqualTo(new SystemEnvironmentOrigin(entry.getKey()));
+			Object actual = replacedMap.get(entry.getKey());
+			assertThat(actual).isEqualTo(entry.getValue());
+			assertThat(replaced.getOrigin(entry.getKey())).isInstanceOf(SystemEnvironmentOrigin.class);
 		}
 	}
 
+	@Test
+	public void replacedPropertySourceWhenPropertyAbsentShouldReturnNullOrigin() throws Exception {
+		SystemEnvironmentPropertySourceEnvironmentPostProcessor postProcessor = new SystemEnvironmentPropertySourceEnvironmentPostProcessor();
+		postProcessor.postProcessEnvironment(this.environment, null);
+		OriginAwareSystemEnvironmentPropertySource replaced = (OriginAwareSystemEnvironmentPropertySource) this.environment.getPropertySources()
+				.get("systemEnvironment");
+		assertThat(replaced.getOrigin("NON_EXISTENT")).isNull();
+	}
+
+	@Test
+	@SuppressWarnings("unchecked")
+	public void replacedPropertySourceShouldResolveProperty() throws Exception {
+		SystemEnvironmentPropertySourceEnvironmentPostProcessor postProcessor = new SystemEnvironmentPropertySourceEnvironmentPostProcessor();
+		Map<String, Object> source = Collections.singletonMap("FOO_BAR_BAZ", "hello");
+		this.environment.getPropertySources().replace("systemEnvironment", new SystemEnvironmentPropertySource("systemEnvironment", source));
+		postProcessor.postProcessEnvironment(this.environment, null);
+		OriginAwareSystemEnvironmentPropertySource replaced = (OriginAwareSystemEnvironmentPropertySource) this.environment.getPropertySources()
+				.get("systemEnvironment");
+		SystemEnvironmentOrigin origin = (SystemEnvironmentOrigin) replaced.getOrigin("foo.bar.baz");
+		assertThat(origin.getProperty()).isEqualTo("FOO_BAR_BAZ");
+		assertThat(replaced.getProperty("foo.bar.baz")).isEqualTo("hello");
+	}
+
 }