interact: implement command state machine

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 {
+func New() *Interact {
-	return &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{}) {
+func (i *Interact) Command(command string, f interface{}) *Command {
-	i.Commands[command] = f
+	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

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 (
+type closePositionTask struct {
-	TestStateStart = 0
+	symbol string
-
+	percentage float64
-	TestStateShowNetAssetValue = 101
+	confirmed bool
-
-	TestStateShowPositionChoosingSymbol = 201
-
-	TestStateClosePositionChoosingSymbol = 301
-	TestStateClosePositionEnteringPercentage = 302
-	TestStateClosePositionConfirming = 303
-)
-
-type TestInteraction struct {
-	State TestState
 }
 
-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)
 }