package tgbot
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"sync"
"testing"
"time"
"github.com/mhsanaei/3x-ui/v3/internal/database"
"github.com/mymmrac/telego"
)
// draftTexts serves the methods the add-client wizard touches and records the
// text of every sendMessage and editMessageText per chat.
func draftTexts(t *testing.T) (string, func(int64) []string) {
t.Helper()
var mu sync.Mutex
texts := map[int64][]string{}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, _ := io.ReadAll(r.Body)
result := any(true)
if r.URL.Path == "/bot"+testBotToken+"/sendMessage" || r.URL.Path == "/bot"+testBotToken+"/editMessageText" {
var payload struct {
ChatID any `json:"chat_id"`
Text string `json:"text"`
}
_ = json.Unmarshal(body, &payload)
chatID := int64(0)
switch v := payload.ChatID.(type) {
case float64:
chatID = int64(v)
}
mu.Lock()
texts[chatID] = append(texts[chatID], payload.Text)
mu.Unlock()
result = map[string]any{"message_id": 1, "date": 0, "chat": map[string]any{"id": chatID, "type": "private"}}
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{"ok": true, "result": result})
}))
t.Cleanup(srv.Close)
return srv.URL, func(chatID int64) []string {
mu.Lock()
defer mu.Unlock()
return append([]string(nil), texts[chatID]...)
}
}
// cardEmail reads the email off a rendered draft card, which is the field the
// wizard assigns when the flow starts.
func cardEmail(t *testing.T, card string) string {
t.Helper()
const marker = "Email: "
start := strings.Index(card, marker)
if start < 0 {
t.Fatalf("not a draft card: %q", card)
}
rest := card[start+len(marker):]
end := strings.Index(rest, "")
if end < 0 {
t.Fatalf("card has an unterminated email: %q", card)
}
return rest[:end]
}
func lastDraftCard(t *testing.T, texts []string) string {
t.Helper()
for i := len(texts) - 1; i >= 0; i-- {
if strings.Contains(texts[i], "Email: ") {
return texts[i]
}
}
t.Fatal("no draft card reached the chat")
return ""
}
// Regression test: one package-level draft per bot meant an admin's new client
// was filled in by another chat's steps.
func TestAddClientDraftIsPerChat(t *testing.T) {
if err := database.InitDB(filepath.Join(t.TempDir(), "x-ui.db")); err != nil {
t.Fatalf("InitDB: %v", err)
}
t.Cleanup(func() { _ = database.CloseDB() })
const (
chatA = int64(7101)
chatB = int64(7202)
)
url, textsFor := draftTexts(t)
swapTestBot(t, url)
origRunning := isRunning
t.Cleanup(func() { isRunning = origRunning })
isRunning = true
callback := func(chatID int64, data string) {
t.Helper()
(&Tgbot{}).answerCallback(&telego.CallbackQuery{
ID: "q1",
From: telego.User{ID: 1},
Data: data,
Message: &telego.Message{MessageID: 7, Chat: telego.Chat{ID: chatID}},
}, true)
}
// Both admins start a client; each card carries the email the wizard just
// generated for that chat.
callback(chatA, "add_client_to 1")
callback(chatB, "add_client_to 2")
emailA := cardEmail(t, lastDraftCard(t, textsFor(chatA)))
emailB := cardEmail(t, lastDraftCard(t, textsFor(chatB)))
if emailA == "" || emailA == emailB {
t.Fatalf("drafts start with the same email %q, want one per chat", emailA)
}
// Chat A renders its card again, with chat B's wizard already past its start.
callback(chatA, "add_client_default_traffic_exp")
if got := cardEmail(t, lastDraftCard(t, textsFor(chatA))); got != emailA {
t.Errorf("chat A's card shows email %q, want its own %q from chat B's draft", got, emailA)
}
if got := cardEmail(t, lastDraftCard(t, textsFor(chatB))); got != emailB {
t.Errorf("chat B's card shows email %q, want %q", got, emailB)
}
}
// Regression test: the draft's lock and map were reached before the admin gate, so
// a report tap queued behind a wizard and any chat a tap came from got stored.
func TestNonWizardCallbackTakesNoDraftLock(t *testing.T) {
const (
heldChat = int64(7303)
spareChat = int64(7404)
)
decliningServer(t)
held := addClientDrafts.forChat(heldChat)
held.Lock()
defer held.Unlock()
tap := func(chatID int64, isAdmin bool, data string) {
(&Tgbot{}).answerCallback(&telego.CallbackQuery{
ID: "q1",
From: telego.User{ID: 1},
Data: data,
Message: &telego.Message{Chat: telego.Chat{ID: chatID}},
}, isAdmin)
}
returns := func(what string, tap func()) {
t.Helper()
done := make(chan struct{})
go func() {
defer close(done)
tap()
}()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatalf("%s waited on the draft lock it never reads", what)
}
}
returns("an admin report tap", func() { tap(heldChat, true, "no_such_admin_action 5") })
returns("a non-admin wizard tap", func() { tap(heldChat, false, "add_client_to 1") })
tap(spareChat, false, "add_client_to 1")
addClientDrafts.mu.Lock()
_, stored := addClientDrafts.drafts[spareChat]
addClientDrafts.mu.Unlock()
if stored {
t.Errorf("draft stored for chat %d, want none until its wizard starts", spareChat)
}
}