From 50ec80217e78d3959472cac9583991939c2c947f Mon Sep 17 00:00:00 2001
From: Albert Salim <albertlimca@gmail.com>
Date: Tue, 11 Oct 2022 20:50:10 +0800
Subject: [PATCH] Add support for Telegram test environment

https://core.telegram.org/bots/webapps#using-bots-in-the-test-environment
---
 api.go      | 14 ++++++++++++--
 api_test.go | 33 +++++++++++++++++++++++++++++++++
 bot.go      |  6 ++++++
 bot_test.go |  2 ++
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/api.go b/api.go
index a66f6211..96b31531 100644
--- a/api.go
+++ b/api.go
@@ -20,7 +20,7 @@ import (
 // It also handles API errors, so you only need to unwrap
 // result field from json data.
 func (b *Bot) Raw(method string, payload interface{}) ([]byte, error) {
-	url := b.URL + "/bot" + b.Token + "/" + method
+	url := b.url(method)
 
 	var buf bytes.Buffer
 	if err := json.NewEncoder(&buf).Encode(payload); err != nil {
@@ -71,6 +71,16 @@ func (b *Bot) Raw(method string, payload interface{}) ([]byte, error) {
 	return data, extractOk(data)
 }
 
+func (b *Bot) url(method string) string {
+	url := b.URL + "/bot" + b.Token
+
+	if b.useTestEnv {
+		url += "/test"
+	}
+
+	return url + "/" + method
+}
+
 func (b *Bot) sendFiles(method string, files map[string]File, params map[string]string) ([]byte, error) {
 	rawFiles := make(map[string]interface{})
 	for name, f := range files {
@@ -116,7 +126,7 @@ func (b *Bot) sendFiles(method string, files map[string]File, params map[string]
 		}
 	}()
 
-	url := b.URL + "/bot" + b.Token + "/" + method
+	url := b.url(method)
 
 	resp, err := b.client.Post(url, writer.FormDataContentType(), pipeReader)
 	if err != nil {
diff --git a/api_test.go b/api_test.go
index d935a6bb..21e2e2a1 100644
--- a/api_test.go
+++ b/api_test.go
@@ -3,6 +3,7 @@ package telebot
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -118,3 +119,35 @@ func TestExtractMessage(t *testing.T) {
 	_, err = extractMessage(data)
 	require.NoError(t, err)
 }
+
+func TestBot_url(t *testing.T) {
+	url := "https://api.telegram.com"
+	token := "my-token"
+	method := "my-method"
+
+	tests := []struct {
+		name        string
+		useTestEnv  bool
+		expectedURL string
+	}{
+		{"when bot is not set to test environment", false, fmt.Sprintf("%s/bot%s/%s", url, token, method)},
+		{"when bot is set to test environment", true, fmt.Sprintf("%s/bot%s/test/%s", url, token, method)},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			settings := Settings{
+				URL:        url,
+				Token:      token,
+				Offline:    true,
+				UseTestEnv: tt.useTestEnv,
+			}
+			bot, err := NewBot(settings)
+
+			assert.NoError(t, err)
+
+			url := bot.url(method)
+
+			assert.Equal(t, tt.expectedURL, url)
+		})
+	}
+}
diff --git a/bot.go b/bot.go
index 16fa608e..7558257b 100644
--- a/bot.go
+++ b/bot.go
@@ -49,6 +49,7 @@ func NewBot(pref Settings) (*Bot, error) {
 		synchronous: pref.Synchronous,
 		verbose:     pref.Verbose,
 		parseMode:   pref.ParseMode,
+		useTestEnv:  pref.UseTestEnv,
 		client:      client,
 	}
 
@@ -80,6 +81,7 @@ type Bot struct {
 	synchronous bool
 	verbose     bool
 	parseMode   ParseMode
+	useTestEnv  bool
 	stop        chan chan struct{}
 	client      *http.Client
 
@@ -122,6 +124,10 @@ type Settings struct {
 
 	// Offline allows to create a bot without network for testing purposes.
 	Offline bool
+
+	// UseTestEnv is used to create a bot on Telegram test environment
+	// https://core.telegram.org/bots/webapps#using-bots-in-the-test-environment
+	UseTestEnv bool
 }
 
 var defaultOnError = func(err error, c Context) {
diff --git a/bot_test.go b/bot_test.go
index cf490c06..9ae7d9e7 100644
--- a/bot_test.go
+++ b/bot_test.go
@@ -67,6 +67,7 @@ func TestNewBot(t *testing.T) {
 	pref.Poller = &LongPoller{Timeout: time.Second}
 	pref.Updates = 50
 	pref.ParseMode = ModeHTML
+	pref.UseTestEnv = true
 	pref.Offline = true
 
 	b, err = NewBot(pref)
@@ -76,6 +77,7 @@ func TestNewBot(t *testing.T) {
 	assert.Equal(t, pref.Poller, b.Poller)
 	assert.Equal(t, 50, cap(b.Updates))
 	assert.Equal(t, ModeHTML, b.parseMode)
+	assert.Equal(t, pref.UseTestEnv, b.useTestEnv)
 }
 
 func TestBotHandle(t *testing.T) {