diff --git a/plugins/babbler/babbler.go b/plugins/babbler/babbler.go index 35f3e22..8e75677 100644 --- a/plugins/babbler/babbler.go +++ b/plugins/babbler/babbler.go @@ -81,16 +81,11 @@ func (p *BabblerPlugin) Message(message msg.Message) bool { numTokens := len(tokens) if numTokens >= 2 && tokens[1] == "says" { - if numTokens > 3 { - p.Bot.SendMessage(message.Channel, "try seabass says [seed-token]") - return true - } - var saying string if len(tokens) == 2 { saying = p.babble(tokens[0]) } else { - saying = p.babbleSeed(tokens[0], tokens[2]) + saying = p.babbleSeed(tokens[0], tokens[2:]) } if saying == "" { p.Bot.SendMessage(message.Channel, "Ze ain't said nothin'") @@ -273,19 +268,37 @@ func getMarkovChain(db *sqlx.DB, who string) (*babbler, error) { } func (p *BabblerPlugin) babble(who string) string { - return p.babbleSeed(who, "") + return p.babbleSeed(who, []string{""}) } -func (p *BabblerPlugin) babbleSeed(who, seed string) string { +func (p *BabblerPlugin) babbleSeed(who string, seed []string) string { if babbler, ok := p.babblers[who]; ok { if len(babbler.start.arcs) == 0 { return "" } - words := []string{seed} + + words := seed var cur *node - if cur, ok = babbler.lookup[seed]; !ok { - return fmt.Sprintf("%s hasn't used the word '%s'", who, seed) + if cur, ok = babbler.lookup[words[0]]; !ok { + if len(words) == 1 { + return fmt.Sprintf("%s hasn't used the word '%s'", who, words[0]) + } else { + return fmt.Sprintf("%s hasn't used the phrase '%s'", who, strings.Join(words, " ")) + } } + + for i := 1; i < len(words); i++ { + if arc, ok := cur.arcs[words[i]]; !ok { + if len(words) == 1 { + return fmt.Sprintf("%s hasn't used the word '%s'", who, words[0]) + } else { + return fmt.Sprintf("%s hasn't used the phrase '%s'", who, strings.Join(words, " ")) + } + } else { + cur = arc.next + } + } + for cur != babbler.end { which := rand.Intn(cur.wordFrequency) sum := 0 diff --git a/plugins/babbler/babbler_test.go b/plugins/babbler/babbler_test.go index be7aa25..7fe4395 100644 --- a/plugins/babbler/babbler_test.go +++ b/plugins/babbler/babbler_test.go @@ -86,6 +86,44 @@ func TestBabblerSeed(t *testing.T) { assert.Contains(t, mb.Messages[0], "long message") } +func TestBabblerMultiSeed(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + c.config.Babbler.DefaultUsers = []string{"seabass"} + assert.NotNil(t, c) + seabass := makeMessage("This is a message") + seabass.User = &user.User{Name: "seabass"} + res := c.Message(seabass) + assert.Len(t, c.babblers, 1) + seabass.Body = "This is another message" + res = c.Message(seabass) + seabass.Body = "This is a long message" + res = c.Message(seabass) + res = c.Message(makeMessage("!seabass says This is a long")) + assert.Len(t, mb.Messages, 1) + assert.True(t, res) + assert.Contains(t, mb.Messages[0], "this is a long message") +} + +func TestBabblerMultiSeed2(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + c.config.Babbler.DefaultUsers = []string{"seabass"} + assert.NotNil(t, c) + seabass := makeMessage("This is a message") + seabass.User = &user.User{Name: "seabass"} + res := c.Message(seabass) + assert.Len(t, c.babblers, 1) + seabass.Body = "This is another message" + res = c.Message(seabass) + seabass.Body = "This is a long message" + res = c.Message(seabass) + res = c.Message(makeMessage("!seabass says is a long")) + assert.Len(t, mb.Messages, 1) + assert.True(t, res) + assert.Contains(t, mb.Messages[0], "is a long message") +} + func TestBabblerBadSeed(t *testing.T) { mb := bot.NewMockBot() c := New(mb) @@ -99,10 +137,29 @@ func TestBabblerBadSeed(t *testing.T) { res = c.Message(seabass) seabass.Body = "This is a long message" res = c.Message(seabass) - res = c.Message(makeMessage("!seabass says long message")) + res = c.Message(makeMessage("!seabass says noooo this is bad")) assert.Len(t, mb.Messages, 1) assert.True(t, res) - assert.Contains(t, mb.Messages[0], "try seabass says [seed-token]") + assert.Contains(t, mb.Messages[0], "seabass hasn't used the phrase 'noooo this is bad'") +} + +func TestBabblerBadSeed2(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + c.config.Babbler.DefaultUsers = []string{"seabass"} + assert.NotNil(t, c) + seabass := makeMessage("This is a message") + seabass.User = &user.User{Name: "seabass"} + res := c.Message(seabass) + assert.Len(t, c.babblers, 1) + seabass.Body = "This is another message" + res = c.Message(seabass) + seabass.Body = "This is a long message" + res = c.Message(seabass) + res = c.Message(makeMessage("!seabass says This is a really")) + assert.Len(t, mb.Messages, 1) + assert.True(t, res) + assert.Contains(t, mb.Messages[0], "seabass hasn't used the phrase 'this is a really'") } func TestBabblerBatch(t *testing.T) {