From a067fe7bbfa53fc0adb8592f20dd073f49685132 Mon Sep 17 00:00:00 2001
From: tijana-ninkovic <tijana.ninkovic@teampicnic.com>
Date: Thu, 18 Jul 2024 10:19:55 +0200
Subject: [PATCH] Group FluxJust rules together

---
 .../refasterrules/ReactorRules.java           | 28 +++++++------------
 .../refasterrules/ReactorRulesTestInput.java  |  8 ++----
 .../refasterrules/ReactorRulesTestOutput.java |  8 ++----
 3 files changed, 14 insertions(+), 30 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 0cf343902be..73675779f12 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
@@ -480,15 +480,20 @@ Flux<T> after() {
   }
 
   /** Prefer {@link Flux#just(Object)} over more contrived alternatives. */
-  static final class FluxJust {
+  static final class FluxJust<T> {
     @BeforeTemplate
-    Flux<Integer> before(int start) {
-      return Flux.range(start, 1);
+    Flux<Integer> before(int value) {
+      return Flux.range(value, 1);
+    }
+
+    @BeforeTemplate
+    Flux<T> before(T value) {
+      return Mono.just(value).repeat().take(1);
     }
 
     @AfterTemplate
-    Flux<Integer> after(int start) {
-      return Flux.just(start);
+    Flux<T> after(T value) {
+      return Flux.just(value);
     }
   }
 
@@ -1727,19 +1732,6 @@ PublisherProbe<T> after() {
     }
   }
 
-  /** Prefer {@link Flux#just(Object)} over more verbose alternatives. */
-  static final class MonoJustRepeatTakeOne<T> {
-    @BeforeTemplate
-    Flux<T> before(T value) {
-      return Mono.just(value).repeat().take(1);
-    }
-
-    @AfterTemplate
-    Flux<T> after(T value) {
-      return Flux.just(value);
-    }
-  }
-
   /** Prefer {@link Mono#as(Function)} when creating a {@link StepVerifier}. */
   static final class StepVerifierFromMono<T> {
     @BeforeTemplate
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 c0cbad05612..63dd6023a9a 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
@@ -182,8 +182,8 @@ ImmutableSet<Flux<?>> testFluxEmpty() {
         Flux.range(0, 0));
   }
 
-  Flux<Integer> testFluxJust() {
-    return Flux.range(0, 1);
+  ImmutableSet<Flux<?>> testFluxJust() {
+    return ImmutableSet.of(Flux.range(0, 1), Mono.just(1).repeat().take(1));
   }
 
   ImmutableSet<Mono<?>> testMonoIdentity() {
@@ -584,10 +584,6 @@ ImmutableSet<PublisherProbe<Void>> testPublisherProbeEmpty() {
     return ImmutableSet.of(PublisherProbe.of(Mono.empty()), PublisherProbe.of(Flux.empty()));
   }
 
-  Flux<Integer> testMonoJustRepeatTakeOne() {
-    return Mono.just(5).repeat().take(1);
-  }
-
   ImmutableSet<StepVerifier.FirstStep<Integer>> testStepVerifierFromMono() {
     return ImmutableSet.of(
         StepVerifier.create(Mono.just(1)), Mono.just(2).flux().as(StepVerifier::create));
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 2106a04b572..e1de8a87676 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
@@ -186,8 +186,8 @@ ImmutableSet<Flux<?>> testFluxEmpty() {
         Flux.empty());
   }
 
-  Flux<Integer> testFluxJust() {
-    return Flux.just(0);
+  ImmutableSet<Flux<?>> testFluxJust() {
+    return ImmutableSet.of(Flux.just(0), Flux.just(1));
   }
 
   ImmutableSet<Mono<?>> testMonoIdentity() {
@@ -573,10 +573,6 @@ ImmutableSet<PublisherProbe<Void>> testPublisherProbeEmpty() {
     return ImmutableSet.of(PublisherProbe.empty(), PublisherProbe.empty());
   }
 
-  Flux<Integer> testMonoJustRepeatTakeOne() {
-    return Flux.just(5);
-  }
-
   ImmutableSet<StepVerifier.FirstStep<Integer>> testStepVerifierFromMono() {
     return ImmutableSet.of(
         Mono.just(1).as(StepVerifier::create), Mono.just(2).as(StepVerifier::create));