From 36444a39c9fff527d60c4e901cbbc463467aeef0 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 f451c94a7f1..1db342f5166 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
@@ -1312,6 +1312,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
@@ -1414,20 +1436,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 410d36fed03..b46bcf55736 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;
@@ -418,6 +419,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();
   }
@@ -450,12 +465,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 62cf1b9c9d0..4793b9b012d 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
@@ -409,6 +409,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();
   }
@@ -439,8 +445,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();
-  }
 }