diff --git a/plugins/riot/src/test/java/com/redis/riot/cli/AbstractRiotTests.java b/plugins/riot/src/test/java/com/redis/riot/cli/AbstractRiotTests.java
index 3791a279f..511accd64 100644
--- a/plugins/riot/src/test/java/com/redis/riot/cli/AbstractRiotTests.java
+++ b/plugins/riot/src/test/java/com/redis/riot/cli/AbstractRiotTests.java
@@ -155,7 +155,8 @@ void teardownTarget() {
 	@BeforeEach
 	void flushAllTarget() throws InterruptedException {
 		targetConnection.sync().flushall();
-		awaitUntil(() -> targetConnection.sync().dbsize().equals(0L));
+		RedisModulesCommands<String, String> sync = targetConnection.sync();
+		awaitEquals(() -> 0L, sync::dbsize);
 	}
 
 	@Override
@@ -421,7 +422,7 @@ void fileImportDump() throws Exception {
 		RedisServerCommands<String, String> sync = connection.sync();
 		sync.flushall();
 		execute("dump-import", this::configureDumpFileImportCommand);
-		Assertions.assertEquals(records.size(), sync.dbsize());
+		awaitEquals(records::size, sync::dbsize);
 	}
 
 	private void configureDumpFileImportCommand(CommandLine.ParseResult parseResult) {
diff --git a/plugins/riot/src/test/java/com/redis/riot/cli/AbstractTestBase.java b/plugins/riot/src/test/java/com/redis/riot/cli/AbstractTestBase.java
index e9534e404..d7c792e43 100644
--- a/plugins/riot/src/test/java/com/redis/riot/cli/AbstractTestBase.java
+++ b/plugins/riot/src/test/java/com/redis/riot/cli/AbstractTestBase.java
@@ -7,6 +7,7 @@
 import java.util.UUID;
 import java.util.concurrent.Callable;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 import org.awaitility.Awaitility;
 import org.codehaus.plexus.util.cli.CommandLineUtils;
@@ -31,6 +32,7 @@
 import org.springframework.core.task.SyncTaskExecutor;
 
 import com.redis.lettucemod.api.StatefulRedisModulesConnection;
+import com.redis.lettucemod.api.sync.RedisModulesCommands;
 import com.redis.lettucemod.util.ClientBuilder;
 import com.redis.lettucemod.util.RedisModulesUtils;
 import com.redis.riot.cli.common.AbstractCommand;
@@ -100,7 +102,8 @@ void setupTest() throws Exception {
 		jobLauncher.setTaskExecutor(new SyncTaskExecutor());
 		stepBuilderFactory = new StepBuilderFactory(jobRepository, new ResourcelessTransactionManager());
 		connection.sync().flushall();
-		awaitUntil(() -> connection.sync().dbsize().equals(0L));
+		RedisModulesCommands<String, String> sync = connection.sync();
+		awaitEquals(() -> 0l, sync::dbsize);
 	}
 
 	protected void awaitUntilFalse(Callable<Boolean> conditionEvaluator) {
@@ -111,6 +114,10 @@ protected void awaitUntil(Callable<Boolean> conditionEvaluator) {
 		Awaitility.await().timeout(DEFAULT_AWAIT_TIMEOUT).until(conditionEvaluator);
 	}
 
+	protected void awaitEquals(Supplier<Object> expected, Supplier<Object> actual) {
+		awaitUntil(() -> expected.get().equals(actual.get()));
+	}
+
 	protected AbstractRedisClient client(RedisServer server) {
 		return ClientBuilder.create(RedisURI.create(server.getRedisURI())).cluster(server.isCluster()).build();
 	}