From d03eae3ddcdea987d49fbb98e8cdd09aa4dfb983 Mon Sep 17 00:00:00 2001
From: Phil Werli <phil.werli@teampicnic.com>
Date: Fri, 18 Nov 2022 15:27:42 +0100
Subject: [PATCH] Extend `MonoFlux` Refaster rule

---
 .../tech/picnic/errorprone/refasterrules/ReactorRules.java    | 2 +-
 .../errorprone/refasterrules/ReactorRulesTestInput.java       | 4 ++--
 .../errorprone/refasterrules/ReactorRulesTestOutput.java      | 4 ++--
 3 files changed, 5 insertions(+), 5 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 382f1b1cde6..0d334758cc6 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
@@ -559,7 +559,7 @@ Flux<S> after(Flux<T> flux) {
   static final class MonoFlux<T> {
     @BeforeTemplate
     Flux<T> before(Mono<T> mono) {
-      return Flux.concat(mono);
+      return Refaster.anyOf(Flux.concat(mono), mono.flatMapMany(Flux::just));
     }
 
     @AfterTemplate
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 b56708a238c..1ff6c359089 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
@@ -178,8 +178,8 @@ ImmutableSet<Flux<Integer>> testFluxMapNotNull() {
         Flux.just(1).switchMap(n -> Mono.fromSupplier(() -> n * 2)));
   }
 
-  Flux<String> testMonoFlux() {
-    return Flux.concat(Mono.just("foo"));
+  ImmutableSet<Flux<String>> testMonoFlux() {
+    return ImmutableSet.of(Flux.concat(Mono.just("foo")), Mono.just("bar").flatMapMany(Flux::just));
   }
 
   ImmutableSet<Mono<Optional<String>>> testMonoCollectToOptional() {
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 03a6788e54a..bb10ccb9bc9 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
@@ -180,8 +180,8 @@ ImmutableSet<Flux<Integer>> testFluxMapNotNull() {
         Flux.just(1).mapNotNull(n -> n * 2));
   }
 
-  Flux<String> testMonoFlux() {
-    return Mono.just("foo").flux();
+  ImmutableSet<Flux<String>> testMonoFlux() {
+    return ImmutableSet.of(Mono.just("foo").flux(), Mono.just("bar").flux());
   }
 
   ImmutableSet<Mono<Optional<String>>> testMonoCollectToOptional() {