From 2935a7db99792e46989d9053e2b15d7727920dcc Mon Sep 17 00:00:00 2001
From: surdeus <surdeus@gmail.com>
Date: Wed, 12 Jul 2023 14:20:52 +0300
Subject: [PATCH] Fixed issues with Update filtering.

---
 src/behx/bot.go      | 12 ++++++++----
 src/behx/context.go  | 15 ++++++++++-----
 src/cmd/test/main.go |  4 +---
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/behx/bot.go b/src/behx/bot.go
index b71ebc0..0df63ee 100644
--- a/src/behx/bot.go
+++ b/src/behx/bot.go
@@ -43,10 +43,11 @@ func (bot *Bot) Run() error {
 	
 	chans := make(map[SessionId] chan *Update)
 	for u := range updates {
+		var sid SessionId
 		if u.Message != nil {	
 			// Create new session if the one does not exist
 			// for this user.
-			sid := SessionId(u.Message.Chat.ID)
+			sid = SessionId(u.Message.Chat.ID)
 			if _, ok := bot.sessions[sid] ; !ok {
 				bot.sessions.Add(sid)
 			} 
@@ -71,10 +72,13 @@ func (bot *Bot) Run() error {
 					continue
 				}
 			}
-			
-			chn := chans[sid]
+		} else if u.CallbackQuery != nil {
+			sid = SessionId(u.CallbackQuery.Message.Chat.ID)
+		}
+		chn, ok := chans[sid]
+		if ok {
 			chn <- &u
-		} 
+		}
 	}
 	
 	return nil
diff --git a/src/behx/context.go b/src/behx/context.go
index 21c7e28..011271d 100644
--- a/src/behx/context.go
+++ b/src/behx/context.go
@@ -1,5 +1,10 @@
 package behx
 
+import (
+	apix "github.com/go-telegram-bot-api/telegram-bot-api/v5"
+	"fmt"
+)
+
 // The type represents way to interact with user in
 // handling functions. Is provided to Act() function always.
 type Context struct {
@@ -13,12 +18,12 @@ func (ctx *Context) handleUpdateChan(updates chan *Update) {
 	session := ctx.Session
 	bot.Start.Act(ctx)
 	for u := range updates {
+		screen := bot.Screens[session.CurrentScreenId]
+		
+		kbd := bot.Keyboards[screen.KeyboardId]
+		btns := kbd.buttonMap()
+		
 		if u.Message != nil {
-			screen := bot.Screens[session.CurrentScreenId]
-			
-			kbd := bot.Keyboards[screen.KeyboardId]
-			btns := kbd.buttonMap()
-			
 			text := u.Message.Text
 			btn, ok := btns[text]
 			
diff --git a/src/cmd/test/main.go b/src/cmd/test/main.go
index d3b2b74..1cd1bbb 100644
--- a/src/cmd/test/main.go
+++ b/src/cmd/test/main.go
@@ -29,9 +29,7 @@ var secondKbd = behx.NewKeyboard(
 	behx.NewButtonRow(
 		behx.NewButton(
 			"❤",
-			behx.NewCustomAction(func(c *behx.Context){
-				log.Println("pressed the button!")
-			}),
+			behx.NewScreenChange("start"),
 		),
 	),
 )