From 5f64fb90ed3500ef5dba025e58aa63c3b3148f24 Mon Sep 17 00:00:00 2001
From: egorbochkarev <bochkarevei@gmail.com>
Date: Fri, 8 Jun 2018 23:12:34 +0300
Subject: [PATCH] Support callback query

---
 .../TelegramLongPollingSessionBot.java        | 37 +++++++++++++------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/telegrambots-chat-session-bot/src/main/java/org/telegram/telegrambots/session/TelegramLongPollingSessionBot.java b/telegrambots-chat-session-bot/src/main/java/org/telegram/telegrambots/session/TelegramLongPollingSessionBot.java
index 908b0350..896a832b 100644
--- a/telegrambots-chat-session-bot/src/main/java/org/telegram/telegrambots/session/TelegramLongPollingSessionBot.java
+++ b/telegrambots-chat-session-bot/src/main/java/org/telegram/telegrambots/session/TelegramLongPollingSessionBot.java
@@ -6,43 +6,56 @@ import org.apache.shiro.session.UnknownSessionException;
 import org.apache.shiro.session.mgt.DefaultSessionManager;
 import org.apache.shiro.session.mgt.SessionContext;
 import org.apache.shiro.session.mgt.eis.AbstractSessionDAO;
+import org.telegram.telegrambots.api.objects.Message;
 import org.telegram.telegrambots.api.objects.Update;
 import org.telegram.telegrambots.bots.TelegramLongPollingBot;
 
 import java.util.Optional;
 
 public abstract class TelegramLongPollingSessionBot extends TelegramLongPollingBot {
-    DefaultSessionManager sessionManager = new DefaultSessionManager();
+    @Setter
+    DefaultSessionManager sessionManager;
 
     @Setter
-    ChatIdConverter chatIdConverter = new DefaultChatIdConverter();
+    ChatIdConverter chatIdConverter;
 
     public TelegramLongPollingSessionBot(){
-        this.setChatIdConverter(new DefaultChatIdConverter());
+        this(new DefaultChatIdConverter());
     }
 
     public TelegramLongPollingSessionBot(ChatIdConverter chatIdConverter){
+        this.setSessionManager(new DefaultSessionManager());
         this.setChatIdConverter(chatIdConverter);
+        AbstractSessionDAO sessionDAO = (AbstractSessionDAO) sessionManager.getSessionDAO();
+        sessionDAO.setSessionIdGenerator(chatIdConverter);
     }
 
     @Override
     public void onUpdateReceived(Update update) {
         Optional<Session> chatSession;
+        Message message;
         if (update.hasMessage()) {
-            chatIdConverter.setSessionId(update.getMessage().getChatId());
-            AbstractSessionDAO sessionDAO = (AbstractSessionDAO) sessionManager.getSessionDAO();
-            sessionDAO.setSessionIdGenerator(chatIdConverter);
-            try {
-                chatSession = Optional.of(sessionManager.getSession(chatIdConverter));
-            } catch (UnknownSessionException e) {
-                SessionContext botSession = new DefaultChatSessionContext(update.getMessage().getChatId(), update.getMessage().getFrom().getUserName());
-                chatSession = Optional.of(sessionManager.start(botSession));
-            }
+            message = update.getMessage();
+        } else if (update.hasCallbackQuery()) {
+            message = update.getCallbackQuery().getMessage();
         } else {
             chatSession = Optional.empty();
+            onUpdateReceived(update, chatSession);
+            return;
         }
+        chatIdConverter.setSessionId(message.getChatId());
+        chatSession = this.getSession(message);
         onUpdateReceived(update, chatSession);
     }
 
+    public Optional<Session> getSession(Message message){
+        try {
+            return Optional.of(sessionManager.getSession(chatIdConverter));
+        } catch (UnknownSessionException e) {
+            SessionContext botSession = new DefaultChatSessionContext(message.getChatId(), message.getFrom().getUserName());
+            return Optional.of(sessionManager.start(botSession));
+        }
+    }
+
     public abstract void onUpdateReceived(Update update, Optional<Session> botSession);
 }