Merge pull request #766 from addo37/ability-replyflow

Fix reply flow registration when using ability definitions

telegrambots-abilities/src/main/java/org/telegram/abilitybots/api/bot/BaseAbilityBot.java

@@ -277,7 +277,8 @@ public abstract class BaseAbilityBot extends DefaultAbsSender implements Ability
 
             // Replies can be standalone or attached to abilities, fetch those too
             Stream<Reply> abilityReplies = abilities.values().stream()
-                    .flatMap(ability -> ability.replies().stream());
+                    .flatMap(ability -> ability.replies().stream())
+                    .flatMap(Reply::stream);
 
             // Now create the replies registry (list)
             replies = Stream.concat(abilityReplies, extensionReplies).collect(

telegrambots-abilities/src/main/java/org/telegram/abilitybots/api/objects/Reply.java

@@ -2,10 +2,8 @@ package org.telegram.abilitybots.api.objects;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Lists;
 import org.telegram.telegrambots.meta.api.objects.Update;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.BiFunction;
@@ -14,7 +12,6 @@ import java.util.function.Predicate;
 import java.util.stream.Stream;
 
 import static com.google.common.collect.Lists.newArrayList;
-import static java.util.Arrays.asList;
 
 /**
  * A reply consists of update conditionals and an action to be applied on the update.
@@ -37,6 +34,13 @@ public class Reply {
     statsEnabled = false;
   }
 
+  Reply(List<Predicate<Update>> conditions, Consumer<Update> action, String name) {
+    this(conditions, action);
+    if (Objects.nonNull(name)) {
+      enableStats(name);
+    }
+  }
+
   public static Reply of(Consumer<Update> action, List<Predicate<Update>> conditions) {
     return new Reply(conditions, action);
   }

telegrambots-abilities/src/main/java/org/telegram/abilitybots/api/objects/ReplyFlow.java

@@ -20,8 +20,8 @@ public class ReplyFlow extends Reply {
 
   private final Set<Reply> nextReplies;
 
-  private ReplyFlow(List<Predicate<Update>> conditions, Consumer<Update> action, Set<Reply> nextReplies) {
+  private ReplyFlow(List<Predicate<Update>> conditions, Consumer<Update> action, Set<Reply> nextReplies, String name) {
-    super(conditions, action);
+    super(conditions, action, name);
     this.nextReplies = nextReplies;
   }
 
@@ -50,6 +50,7 @@ public class ReplyFlow extends Reply {
     private List<Predicate<Update>> conds;
     private Consumer<Update> action;
     private Set<Reply> nextReplies;
+    private String name;
 
     private ReplyFlowBuilder(DBContext db, int id) {
       conds = new ArrayList<>();
@@ -67,6 +68,11 @@ public class ReplyFlow extends Reply {
       return this;
     }
 
+    public ReplyFlowBuilder enableStats(String name) {
+      this.name = name;
+      return this;
+    }
+
     public ReplyFlowBuilder onlyIf(Predicate<Update> pred) {
       conds.add(pred);
       return this;
@@ -79,7 +85,7 @@ public class ReplyFlow extends Reply {
         db.<Long, Integer>getMap(STATES).remove(chatId);
       });
 
-      Reply statefulReply = Reply.of(statefulAction, statefulConditions);
+      Reply statefulReply = new Reply(statefulConditions, statefulAction, nextReply.name());
       nextReplies.add(statefulReply);
       return this;
     }
@@ -87,7 +93,7 @@ public class ReplyFlow extends Reply {
     public ReplyFlowBuilder next(ReplyFlow nextReplyFlow) {
       List<Predicate<Update>> statefulConditions = toStateful(nextReplyFlow.conditions());
 
-      ReplyFlow statefulReplyFlow = new ReplyFlow(statefulConditions, nextReplyFlow.action(), nextReplyFlow.nextReplies());
+      ReplyFlow statefulReplyFlow = new ReplyFlow(statefulConditions, nextReplyFlow.action(), nextReplyFlow.nextReplies(), nextReplyFlow.name());
       nextReplies.add(statefulReplyFlow);
       return this;
     }
@@ -95,12 +101,21 @@ public class ReplyFlow extends Reply {
     public ReplyFlow build() {
       if (action == null)
         action = upd -> {};
-      Consumer<Update> statefulAction = action.andThen(upd -> {
+
+      Consumer<Update> statefulAction;
+      if (nextReplies.size() > 0) {
+        statefulAction = action.andThen(upd -> {
           Long chatId = AbilityUtils.getChatId(upd);
           db.<Long, Integer>getMap(STATES).put(chatId, id);
         });
+      } else {
+        statefulAction = action.andThen(upd -> {
+          Long chatId = AbilityUtils.getChatId(upd);
+          db.<Long, Integer>getMap(STATES).remove(chatId);
+        });
+      }
 
-      return new ReplyFlow(conds, statefulAction, nextReplies);
+      return new ReplyFlow(conds, statefulAction, nextReplies, name);
     }
 
     @NotNull

telegrambots-abilities/src/test/java/org/telegram/abilitybots/api/bot/ReplyFlowTest.java

@@ -5,17 +5,18 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.telegram.abilitybots.api.db.DBContext;
-import org.telegram.abilitybots.api.objects.Flag;
+import org.telegram.abilitybots.api.objects.*;
-import org.telegram.abilitybots.api.objects.Reply;
-import org.telegram.abilitybots.api.objects.ReplyFlow;
 import org.telegram.abilitybots.api.sender.MessageSender;
 import org.telegram.abilitybots.api.sender.SilentSender;
 import org.telegram.telegrambots.meta.api.objects.Update;
 import org.telegram.telegrambots.meta.api.objects.polls.Poll;
 
 import java.io.IOException;
+import java.util.Set;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
+import static com.google.common.collect.Sets.newHashSet;
 import static org.apache.commons.lang3.StringUtils.EMPTY;
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;
@@ -122,6 +123,35 @@ public class ReplyFlowTest {
     assertTrue(bot.filterReply(update));
   }
 
+  @Test
+  void replyFlowsAreWorkingWhenDefinedInAbilities() {
+    Update update1 = mockFullUpdate(bot, USER, "one");
+    Update update2 = mockFullUpdate(bot, USER, "two");
+    long chatId = getChatId(update1);
+
+    // Trigger and verify first reply stage
+    assertFalse(bot.filterReply(update1));
+
+    verify(silent, only()).send("First reply", chatId);
+    assertTrue(db.<Long, Integer>getMap(STATES).containsKey(chatId), "User is not in initial state");
+    // Resetting the mock now helps with verification later
+    reset(silent);
+
+    // Trigger and verify second reply stage
+    assertFalse(bot.filterReply(update2));
+
+    verify(silent, only()).send("Second reply", chatId);
+    assertFalse(db.<Long, Integer>getMap(STATES).containsKey(chatId), "User is still in a state");
+  }
+
+  @Test
+  void replyFlowsPertainNames() {
+    Update update1 = mockFullUpdate(bot, USER, "one");
+    long chatId = getChatId(update1);
+    Set<String> replyNames = bot.replies().stream().map(Reply::name).collect(Collectors.toSet());
+    replyNames.containsAll(newHashSet("FIRST", "SECOND"));
+  }
+
   public static class ReplyFlowBot extends AbilityBot {
 
     private ReplyFlowBot(String botToken, String botUsername, DBContext db) {
@@ -153,6 +183,33 @@ public class ReplyFlowTest {
           .build();
     }
 
+    public Ability replyFlowsWithAbility() {
+      Reply replyWithVk = ReplyFlow.builder(db, 2)
+          .enableStats("SECOND")
+          .action(upd -> {
+            silent.send("Second reply", upd.getMessage().getChatId());
+          })
+          .onlyIf(hasMessageWith("two"))
+          .build();
+
+      Reply replyWithNickname = ReplyFlow.builder(db, 1)
+          .enableStats("FIRST")
+          .action(upd -> {
+            silent.send("First reply", upd.getMessage().getChatId());
+          })
+          .onlyIf(hasMessageWith("one"))
+          .next(replyWithVk)
+          .build();
+
+      return Ability.builder()
+          .name("trigger")
+          .privacy(Privacy.PUBLIC)
+          .locality(Locality.ALL)
+          .action(ctx -> silent.send("I'm in an ability", ctx.chatId()))
+          .reply(replyWithNickname)
+          .build();
+    }
+
     @NotNull
     private Predicate<Update> hasMessageWith(String msg) {
       return upd -> Flag.MESSAGE.test(upd) && upd.getMessage().getText().equalsIgnoreCase(msg);