From a350afa3cff4e6df1a68f394250affc39620a4e7 Mon Sep 17 00:00:00 2001
From: Johnny Lim <izeye@naver.com>
Date: Thu, 23 Mar 2023 08:07:05 +0900
Subject: [PATCH] Add supplier variant for Observation.Context.getOrDefault()
 (#3708)

---
 .../micrometer/observation/Observation.java   | 13 +++++++++
 .../observation/ObservationContextTest.java   | 27 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/micrometer-observation/src/main/java/io/micrometer/observation/Observation.java b/micrometer-observation/src/main/java/io/micrometer/observation/Observation.java
index 126757ac7f..fa6090ef0a 100644
--- a/micrometer-observation/src/main/java/io/micrometer/observation/Observation.java
+++ b/micrometer-observation/src/main/java/io/micrometer/observation/Observation.java
@@ -1294,6 +1294,19 @@ interface ContextView {
          */
         <T> T getOrDefault(Object key, T defaultObject);
 
+        /**
+         * Returns an element or default if not present.
+         * @param key key
+         * @param defaultObjectSupplier supplier for default object to return
+         * @param <T> value type
+         * @return object or default if not present
+         * @since 1.11.0
+         */
+        default <T> T getOrDefault(Object key, Supplier<T> defaultObjectSupplier) {
+            T value = get(key);
+            return value != null ? value : defaultObjectSupplier.get();
+        }
+
         /**
          * Returns low cardinality key values.
          * @return low cardinality key values
diff --git a/micrometer-observation/src/test/java/io/micrometer/observation/ObservationContextTest.java b/micrometer-observation/src/test/java/io/micrometer/observation/ObservationContextTest.java
index b67848db8d..7af36882c4 100644
--- a/micrometer-observation/src/test/java/io/micrometer/observation/ObservationContextTest.java
+++ b/micrometer-observation/src/test/java/io/micrometer/observation/ObservationContextTest.java
@@ -21,7 +21,10 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.util.function.Supplier;
+
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.*;
 
 /**
  * Tests for {@link Observation.Context}.
@@ -68,6 +71,30 @@ void getOrDefaultShouldUseFallbackValue() {
         assertThat(context.getOrDefault(Integer.class, 123)).isEqualTo(123);
     }
 
+    @Test
+    void getOrDefaultSupplierWhenKeyIsPresent() {
+        context.put(String.class, "42");
+
+        @SuppressWarnings("unchecked")
+        Supplier<String> defaultSupplier = mock(Supplier.class);
+        when(defaultSupplier.get()).thenReturn("abc");
+
+        assertThat(context.getOrDefault(String.class, defaultSupplier)).isEqualTo("42");
+        verifyNoInteractions(defaultSupplier);
+    }
+
+    @Test
+    void getOrDefaultSupplierWhenKeyIsMissing() {
+        context.put(String.class, "42");
+
+        @SuppressWarnings("unchecked")
+        Supplier<Integer> defaultSupplier = mock(Supplier.class);
+        when(defaultSupplier.get()).thenReturn(123);
+
+        assertThat(context.getOrDefault(Integer.class, defaultSupplier)).isEqualTo(123);
+        verify(defaultSupplier, times(1)).get();
+    }
+
     @Test
     void getRequiredShouldFailIfThereIsNoValue() {
         context.put(String.class, "42");