From 2e5d1f1e8726b770228f1254ba57b4fc054eecb6 Mon Sep 17 00:00:00 2001
From: tijana-ninkovic <tijana.ninkovic@teampicnic.com>
Date: Sat, 20 Jul 2024 12:57:45 +0200
Subject: [PATCH] Extend `FluxJust` Refaster rule (#1253)

---
 .../errorprone/refasterrules/ReactorRules.java    | 15 ++++++++++-----
 .../refasterrules/ReactorRulesTestInput.java      |  4 ++--
 .../refasterrules/ReactorRulesTestOutput.java     |  4 ++--
 3 files changed, 14 insertions(+), 9 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 273e0eea17..1f97c40fd3 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
@@ -481,15 +481,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);
     }
   }
 
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 1efc888cc4..3702be49e9 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
@@ -183,8 +183,8 @@ ImmutableSet<Flux<?>> testFluxEmpty() {
         Flux.range(0, 0));
   }
 
-  Flux<Integer> testFluxJust() {
-    return Flux.range(0, 1);
+  ImmutableSet<Flux<Integer>> testFluxJust() {
+    return ImmutableSet.of(Flux.range(0, 1), Mono.just(2).repeat().take(1));
   }
 
   ImmutableSet<Mono<?>> testMonoIdentity() {
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 7e4c325bf1..7cb8a551b2 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
@@ -187,8 +187,8 @@ ImmutableSet<Flux<?>> testFluxEmpty() {
         Flux.empty());
   }
 
-  Flux<Integer> testFluxJust() {
-    return Flux.just(0);
+  ImmutableSet<Flux<Integer>> testFluxJust() {
+    return ImmutableSet.of(Flux.just(0), Flux.just(2));
   }
 
   ImmutableSet<Mono<?>> testMonoIdentity() {