From 3cc11badac4a91b93f0e0409d12c35c24835a3d7 Mon Sep 17 00:00:00 2001 From: c9s Date: Thu, 13 Jan 2022 00:59:58 +0800 Subject: [PATCH] interact: implement command state machine --- pkg/interact/interact.go | 133 ++++++++++++++++++++++++++++++++-- pkg/interact/interact_test.go | 72 +++++++++++++----- 2 files changed, 180 insertions(+), 25 deletions(-) diff --git a/pkg/interact/interact.go b/pkg/interact/interact.go index e7833108c..b9b5eaf19 100644 --- a/pkg/interact/interact.go +++ b/pkg/interact/interact.go @@ -1,6 +1,7 @@ package interact import ( + "fmt" "reflect" "strconv" "strings" @@ -9,17 +10,135 @@ import ( "gopkg.in/tucnak/telebot.v2" ) +type Command struct { + // Name is the command name + Name string + + // StateF is the command handler function + F interface{} + + stateID int + states map[string]string + statesFunc map[string]interface{} + initState string +} + +func NewCommand(name string, f interface{}) *Command { + c := &Command{ + Name: name, + F: f, + states: make(map[string]string), + statesFunc: make(map[string]interface{}), + initState: name + "_" + strconv.Itoa(0), + } + return c.Next(f) +} + +func (c *Command) Next(f interface{}) *Command { + curState := c.Name + "_" + strconv.Itoa(c.stateID) + c.stateID++ + nextState := c.Name + "_" + strconv.Itoa(c.stateID) + + c.states[curState] = nextState + c.statesFunc[curState] = f + return c +} + // Interact implements the interaction between bot and message software. type Interact struct { - Commands map[string]interface{} + commands map[string]*Command + + states map[string]string + statesFunc map[string]interface{} + curState string } -func New() *Interact { - return &Interact{} +func New() *Interact { + return &Interact{ + commands: make(map[string]*Command), + states: make(map[string]string), + statesFunc: make(map[string]interface{}), + } } -func (i *Interact) Command(command string, f interface{}) { - i.Commands[command] = f +func (i *Interact) Command(command string, f interface{}) *Command { + cmd := NewCommand(command, f) + i.commands[command] = cmd + return cmd +} + +func (i *Interact) getNextState(currentState string) (nextState string, end bool) { + var ok bool + nextState, ok = i.states[currentState] + if ok { + end = false + return nextState, end + } + + end = true + return nextState, end +} + +func (i *Interact) handleResponse(text string) error { + args := parseCommand(text) + + f, ok := i.statesFunc[i.curState] + if !ok { + return fmt.Errorf("state function of %s is not defined", i.curState) + } + + err := parseFuncArgsAndCall(f, args) + if err != nil { + return err + } + + nextState, end := i.getNextState(i.curState) + if end { + return nil + } + + i.curState = nextState + return nil +} + +func (i *Interact) runCommand(command string, args ...string) error { + cmd, ok := i.commands[command] + if !ok { + return fmt.Errorf("command %s not found", command) + } + + i.curState = cmd.initState + err := parseFuncArgsAndCall(cmd.F, args) + if err != nil { + return err + } + + // if we can successfully execute the command, then we can go to the next state. + nextState, end := i.getNextState(i.curState) + if end { + return nil + } + + i.curState = nextState + return nil +} + +func (i *Interact) init() error { + for n, cmd := range i.commands { + _ = n + for s1, s2 := range cmd.states { + if _, exist := i.states[s1]; exist { + return fmt.Errorf("state %s already exists", s1) + } + + i.states[s1] = s2 + } + for s, f := range cmd.statesFunc { + i.statesFunc[s] = f + } + } + + return nil } func (i *Interact) HandleTelegramMessage(msg *telebot.Message) { @@ -36,10 +155,10 @@ func parseCommand(src string) (args []string) { s.Filename = "command" for tok := s.Scan(); tok != scanner.EOF; tok = s.Scan() { text := s.TokenText() - if text[0] == '"' && text[len(text) - 1] == '"' { + if text[0] == '"' && text[len(text)-1] == '"' { text, _ = strconv.Unquote(text) } - args = append(args,text) + args = append(args, text) } return args diff --git a/pkg/interact/interact_test.go b/pkg/interact/interact_test.go index e4c5b0055..dd0cf05cf 100644 --- a/pkg/interact/interact_test.go +++ b/pkg/interact/interact_test.go @@ -43,34 +43,70 @@ func Test_parseCommand(t *testing.T) { assert.Equal(t, "market", args[3]) } -type TestState int -const ( - TestStateStart = 0 - - TestStateShowNetAssetValue = 101 - - TestStateShowPositionChoosingSymbol = 201 - - TestStateClosePositionChoosingSymbol = 301 - TestStateClosePositionEnteringPercentage = 302 - TestStateClosePositionConfirming = 303 -) - -type TestInteraction struct { - State TestState +type closePositionTask struct { + symbol string + percentage float64 + confirmed bool } -func (m *TestInteraction) HandleMessage() { - +type TestInteraction struct { + closePositionTask closePositionTask } func (m *TestInteraction) Commands(interact *Interact) { - interact.Command("closePosition", func() { + interact.Command("closePosition", func() error { + // send symbol options + return nil + }).Next(func(symbol string) error { + // get symbol from user + m.closePositionTask.symbol = symbol + // send percentage options + return nil + }).Next(func(percentage float64) error { + // get percentage from user + m.closePositionTask.percentage = percentage + + // send confirmation + return nil + }).Next(func(confirmed bool) error { + m.closePositionTask.confirmed = confirmed + // call position close + + // reply result + return nil }) } func TestCustomInteraction(t *testing.T) { + globalInteraction := New() + testInteraction := &TestInteraction{} + testInteraction.Commands(globalInteraction) + err := globalInteraction.init() + assert.NoError(t, err) + + err = globalInteraction.runCommand("closePosition") + assert.NoError(t, err) + + assert.Equal(t, "closePosition_1", globalInteraction.curState) + + err = globalInteraction.handleResponse("BTCUSDT") + assert.NoError(t, err) + assert.Equal(t, "closePosition_2", globalInteraction.curState) + + err = globalInteraction.handleResponse("0.20") + assert.NoError(t, err) + assert.Equal(t, "closePosition_3", globalInteraction.curState) + + err = globalInteraction.handleResponse("true") + assert.NoError(t, err) + assert.Equal(t, "closePosition_4", globalInteraction.curState) + + assert.Equal(t, closePositionTask{ + symbol: "BTCUSDT", + percentage: 0.2, + confirmed: true, + }, testInteraction.closePositionTask) }