Merge pull request #44 from velour/babbler_multi_seed

allow the babbler to be seeded with a full phrase

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

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) {