From f02ed8751e9f2fb866b10118c7f2ac530c2b0cfe Mon Sep 17 00:00:00 2001
From: Yordan Grancharov <phrone@gmail.com>
Date: Sun, 21 Jul 2024 14:49:13 +0300
Subject: [PATCH] Add test cases for computeAsync, improve fix

Signed-off-by: Yordan Grancharov <phrone@gmail.com>
---
 .../main/java/org/redisson/RedissonMap.java   | 11 ++++---
 .../test/java/org/redisson/BaseMapTest.java   | 29 +++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/redisson/src/main/java/org/redisson/RedissonMap.java b/redisson/src/main/java/org/redisson/RedissonMap.java
index 3fffcd3bc..2c14aab82 100644
--- a/redisson/src/main/java/org/redisson/RedissonMap.java
+++ b/redisson/src/main/java/org/redisson/RedissonMap.java
@@ -258,11 +258,14 @@ public class RedissonMap<K, V> extends RedissonExpirable implements RMap<K, V> {
                 return oldValueFuture.thenCompose(oldValue -> {
                     return CompletableFuture.supplyAsync(() -> remappingFunction.apply(key, oldValue), getServiceManager().getExecutor())
                             .thenCompose(newValue -> {
-                                if (newValue != null) {
-                                    return fastPutAsync(key, newValue)
-                                            .thenApply(rr -> newValue);
+                                if (newValue == null) {
+                                    if (oldValue != null) {
+                                        return fastRemoveAsync(key)
+                                                .thenApply(rr -> newValue);
+                                    }
+                                    return CompletableFuture.completedFuture(newValue);
                                 }
-                                return fastRemoveAsync(key)
+                                return fastPutAsync(key, newValue)
                                         .thenApply(rr -> newValue);
                             });
                 }).whenComplete((c, e) -> {
diff --git a/redisson/src/test/java/org/redisson/BaseMapTest.java b/redisson/src/test/java/org/redisson/BaseMapTest.java
index 5d4c0ffd8..fe183d35c 100644
--- a/redisson/src/test/java/org/redisson/BaseMapTest.java
+++ b/redisson/src/test/java/org/redisson/BaseMapTest.java
@@ -269,14 +269,43 @@ public abstract class BaseMapTest extends RedisDockerTest {
     @Test
     public void testCompute() {
         RMap<String, String> map = getMap("map");
+
         map.compute("1", (key, oldValue) -> {
             return "12";
         });
         assertThat(map.get("1")).isEqualTo("12");
 
         map.compute("1", (key, oldValue) -> {
+            return (oldValue == null) ? "12" : oldValue.concat("34");
+        });
+        assertThat(map.get("1")).isEqualTo("1234");
+
+        map.compute("1", (key, oldValue) -> {
+            return null;
+        });
+        assertThat(map.get("1")).isNull();
+    }
+
+    @Test
+    public void testComputeAsync() {
+        RMap<String, String> map = getMap("mapAsync");
+
+        RFuture<String> res1 = map.computeAsync("1", (key, oldValue) -> {
+            return "12";
+        });
+        assertThat(res1.toCompletableFuture().join()).isEqualTo("12");
+        assertThat(map.get("1")).isEqualTo("12");
+
+        RFuture<String> res2 = map.computeAsync("1", (key, oldValue) -> {
+            return (oldValue == null) ? "12" : oldValue.concat("34");
+        });
+        assertThat(res2.toCompletableFuture().join()).isEqualTo("1234");
+        assertThat(map.get("1")).isEqualTo("1234");
+
+        RFuture<String> res3 = map.computeAsync("1", (key, oldValue) -> {
             return null;
         });
+        assertThat(res3.toCompletableFuture().join()).isNull();
         assertThat(map.get("1")).isNull();
     }