From 200aafc3dd0fb0d90d6b3094f8f985f551c4d90d Mon Sep 17 00:00:00 2001
From: Rick Ossendrijver <rick.ossendrijver@gmail.com>
Date: Tue, 9 May 2023 14:25:33 +0200
Subject: [PATCH] Suggestions

---
 .../refasterrules/ReactorRules.java           | 38 +++++++++++--------
 .../refasterrules/ReactorRulesTestInput.java  | 23 +++++++----
 .../refasterrules/ReactorRulesTestOutput.java | 10 +++--
 3 files changed, 43 insertions(+), 28 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
index 8ef3a658b73..eff0042e706 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
@@ -1299,6 +1299,28 @@ StepVerifier.Step<T> after(StepVerifier.Step<T> step, T object) {
     }
   }
 
+  /** Avoid collecting when verifying only the given element is in the given {@link Flux}. */
+  static final class FluxAsStepVerifierExpectNext<T> {
+    @BeforeTemplate
+    StepVerifier.Step<ImmutableList<T>> before(Flux<T> flux, T object) {
+      return flux.collect(toImmutableList())
+          .as(StepVerifier::create)
+          .assertNext(list -> assertThat(list).containsExactly(object));
+    }
+
+    @BeforeTemplate
+    StepVerifier.Step<ImmutableSet<T>> before2(Flux<T> flux, T object) {
+      return flux.collect(toImmutableSet())
+          .as(StepVerifier::create)
+          .assertNext(set -> assertThat(set).containsExactly(object));
+    }
+
+    @AfterTemplate
+    StepVerifier.Step<T> after(Flux<T> flux, T object) {
+      return flux.as(StepVerifier::create).expectNext(object);
+    }
+  }
+
   /** Prefer {@link StepVerifier.LastStep#verifyComplete()} over more verbose alternatives. */
   static final class StepVerifierLastStepVerifyComplete {
     @BeforeTemplate
@@ -1401,20 +1423,4 @@ Duration after(StepVerifier.LastStep step, Duration duration) {
       return step.verifyTimeout(duration);
     }
   }
-
-  /** Avoid collecting when verifying only the given element is in the given {@link Flux}. */
-  static final class VerifyOnlyElementInFlux<T> {
-    @BeforeTemplate
-    Duration before(Flux<T> flux, T object) {
-      return flux.collect(toImmutableList())
-          .as(StepVerifier::create)
-          .assertNext(list -> assertThat(list).containsExactly(object))
-          .verifyComplete();
-    }
-
-    @AfterTemplate
-    Duration after(Flux<T> flux, T object) {
-      return flux.as(StepVerifier::create).expectNext(object).verifyComplete();
-    }
-  }
 }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
index 8ee932d952a..24dd674302d 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
@@ -1,6 +1,7 @@
 package tech.picnic.errorprone.refasterrules;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.MoreCollectors.toOptional;
 import static java.util.Comparator.reverseOrder;
 import static java.util.function.Function.identity;
@@ -414,6 +415,20 @@ ImmutableSet<StepVerifier.Step<String>> testStepVerifierStepExpectNext() {
         StepVerifier.create(Mono.just("baz")).expectNextMatches("qux"::equals));
   }
 
+  ImmutableSet<Duration> testFluxAsStepVerifierExpectNext() {
+    return ImmutableSet.of(
+        Flux.just(1)
+            .collect(toImmutableList())
+            .as(StepVerifier::create)
+            .assertNext(list -> assertThat(list).containsExactly(2))
+            .verifyComplete(),
+        Flux.just(3)
+            .collect(toImmutableSet())
+            .as(StepVerifier::create)
+            .assertNext(set -> assertThat(set).containsExactly(4))
+            .verifyComplete());
+  }
+
   Duration testStepVerifierLastStepVerifyComplete() {
     return StepVerifier.create(Mono.empty()).expectComplete().verify();
   }
@@ -446,12 +461,4 @@ Duration testStepVerifierLastStepVerifyErrorMessage() {
   Duration testStepVerifierLastStepVerifyTimeout() {
     return StepVerifier.create(Mono.empty()).expectTimeout(Duration.ZERO).verify();
   }
-
-  Duration testVerifyOnlyElementInFlux() {
-    return Flux.just(1)
-        .collect(toImmutableList())
-        .as(StepVerifier::create)
-        .assertNext(list -> assertThat(list).containsExactly(1))
-        .verifyComplete();
-  }
 }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
index 353b67184ec..a46bb8c7914 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
@@ -405,6 +405,12 @@ ImmutableSet<StepVerifier.Step<String>> testStepVerifierStepExpectNext() {
         StepVerifier.create(Mono.just("baz")).expectNext("qux"));
   }
 
+  ImmutableSet<Duration> testFluxAsStepVerifierExpectNext() {
+    return ImmutableSet.of(
+        Flux.just(1).as(StepVerifier::create).expectNext(2).verifyComplete(),
+        Flux.just(3).as(StepVerifier::create).expectNext(4).verifyComplete());
+  }
+
   Duration testStepVerifierLastStepVerifyComplete() {
     return StepVerifier.create(Mono.empty()).verifyComplete();
   }
@@ -435,8 +441,4 @@ Duration testStepVerifierLastStepVerifyErrorMessage() {
   Duration testStepVerifierLastStepVerifyTimeout() {
     return StepVerifier.create(Mono.empty()).verifyTimeout(Duration.ZERO);
   }
-
-  Duration testVerifyOnlyElementInFlux() {
-    return Flux.just(1).as(StepVerifier::create).expectNext(1).verifyComplete();
-  }
 }