From e2a919b8fa791ed49943851b11d7db7365cb3fee Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@vmware.com>
Date: Thu, 7 Jul 2022 14:49:48 +0200
Subject: [PATCH] Added assertions for handled contexts; fixes gh-3274

---
 .../tck/TestObservationRegistryAssert.java    | 23 +++++++++++++++++--
 .../TestObservationRegistryAssertTests.java   | 18 +++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/micrometer-observation-test/src/main/java/io/micrometer/observation/tck/TestObservationRegistryAssert.java b/micrometer-observation-test/src/main/java/io/micrometer/observation/tck/TestObservationRegistryAssert.java
index ea31d2a0db..19e0e422dc 100644
--- a/micrometer-observation-test/src/main/java/io/micrometer/observation/tck/TestObservationRegistryAssert.java
+++ b/micrometer-observation-test/src/main/java/io/micrometer/observation/tck/TestObservationRegistryAssert.java
@@ -16,6 +16,7 @@
 package io.micrometer.observation.tck;
 
 import io.micrometer.observation.Observation;
+import org.assertj.core.api.ThrowingConsumer;
 
 import java.util.List;
 import java.util.Objects;
@@ -104,8 +105,8 @@ private String observationNames(List<TestObservationRegistry.TestObservationCont
     }
 
     /**
-     * Verifies that there's only one {@link Observation} with a given name (ignoring
-     * case) and continues assertions for it.
+     * Verifies that there's at least one {@link Observation} with a given name (ignoring
+     * case) and continues assertions for the first found one.
      * @return this
      * @throws AssertionError if there is no matching observation
      */
@@ -125,6 +126,11 @@ public That hasObservationWithNameEqualToIgnoringCase(String name) {
         return new That(testObservationContext, this);
     }
 
+    /**
+     * Verifies that there are no observations registered.
+     * @return this
+     * @throws AssertionError if there are any registered observations
+     */
     public void doesNotHaveAnyObservation() {
         List<TestObservationRegistry.TestObservationContext> contexts = this.actual.getContexts();
         if (!contexts.isEmpty()) {
@@ -133,6 +139,19 @@ public void doesNotHaveAnyObservation() {
         }
     }
 
+    /**
+     * Verifies that all handled contexts satisfy the provided lambda.
+     * @param contextConsumer lambda to assert all handled contexts
+     * @return this
+     */
+    public TestObservationRegistryAssert hasHandledContextsThatSatisfy(
+            ThrowingConsumer<List<Observation.Context>> contextConsumer) {
+        isNotNull();
+        contextConsumer.accept(actual.getContexts().stream()
+                .map(TestObservationRegistry.TestObservationContext::getContext).collect(Collectors.toList()));
+        return this;
+    }
+
     /**
      * Provides assertions for {@link Observation} and allows coming back to
      * {@link TestObservationRegistryAssert}.
diff --git a/micrometer-observation-test/src/test/java/io/micrometer/observation/tck/TestObservationRegistryAssertTests.java b/micrometer-observation-test/src/test/java/io/micrometer/observation/tck/TestObservationRegistryAssertTests.java
index fc105acc12..e78f799f55 100644
--- a/micrometer-observation-test/src/test/java/io/micrometer/observation/tck/TestObservationRegistryAssertTests.java
+++ b/micrometer-observation-test/src/test/java/io/micrometer/observation/tck/TestObservationRegistryAssertTests.java
@@ -17,6 +17,7 @@
 
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
+import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 import static io.micrometer.observation.tck.TestObservationRegistryAssert.assertThat;
@@ -111,6 +112,23 @@ void should_not_fail_when_observation_with_name_ignoring_case_found() {
                 .hasObservationWithNameEqualToIgnoringCase("foo").that().hasBeenStarted());
     }
 
+    @Test
+    void should_fail_when_no_contexts_satisfy_the_assertion() {
+        Observation.createNotStarted("foo", registry).start().stop();
+
+        thenThrownBy(() -> TestObservationRegistryAssert.assertThat(registry)
+                .hasHandledContextsThatSatisfy(contexts -> Assertions.assertThat(contexts).hasSize(2)))
+                        .isInstanceOf(AssertionError.class);
+    }
+
+    @Test
+    void should_not_fail_when_contexts_satisfy_the_assertions() {
+        Observation.createNotStarted("FOO", registry).start().stop();
+
+        thenNoException().isThrownBy(() -> TestObservationRegistryAssert.assertThat(registry)
+                .hasHandledContextsThatSatisfy(contexts -> Assertions.assertThat(contexts).hasSize(1)));
+    }
+
     @Test
     void should_jump_to_and_back_from_context_assert() {
         new Example(registry).run();