From a86ef4bebdd83056f3a18cbf431d8d38a7a89b53 Mon Sep 17 00:00:00 2001
From: Scott Kiesel <skiesel@Scotts-MacBook-Pro.local>
Date: Wed, 19 Oct 2016 19:34:16 -0400
Subject: [PATCH] stk5's wish come true

---
 plugins/babbler/babbler.go      | 42 +++++++++++++++++------
 plugins/babbler/babbler_test.go | 60 +++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 11 deletions(-)

diff --git a/plugins/babbler/babbler.go b/plugins/babbler/babbler.go
index 01d97af..35f3e22 100644
--- a/plugins/babbler/babbler.go
+++ b/plugins/babbler/babbler.go
@@ -78,13 +78,25 @@ func (p *BabblerPlugin) Message(message msg.Message) bool {
 
 	lowercase := strings.ToLower(message.Body)
 	tokens := strings.Fields(lowercase)
+	numTokens := len(tokens)
 
-	if len(tokens) == 2 && tokens[1] == "says" {
-		saying := p.babble(tokens[0])
+	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])
+		}
 		if saying == "" {
 			p.Bot.SendMessage(message.Channel, "Ze ain't said nothin'")
+		} else {
+			p.Bot.SendMessage(message.Channel, saying)
 		}
-		p.Bot.SendMessage(message.Channel, saying)
 		return true
 	} else if len(tokens) == 4 && strings.Index(lowercase, "initialize babbler for ") == 0 {
 		who := tokens[3]
@@ -223,16 +235,17 @@ func addToMarkovChain(babble *babbler, phrase string) {
 }
 
 func newBabbler() *babbler {
+	start := &node{
+		wordFrequency: 0,
+		arcs:          map[string]*arc{},
+	}
 	return &babbler{
-		start: &node{
-			wordFrequency: 0,
-			arcs:          map[string]*arc{},
-		},
+		start: start,
 		end: &node{
 			wordFrequency: 0,
 			arcs:          map[string]*arc{},
 		},
-		lookup: map[string]*node{},
+		lookup: map[string]*node{"": start},
 	}
 }
 
@@ -260,12 +273,19 @@ func getMarkovChain(db *sqlx.DB, who string) (*babbler, error) {
 }
 
 func (p *BabblerPlugin) babble(who string) string {
+	return p.babbleSeed(who, "")
+}
+
+func (p *BabblerPlugin) babbleSeed(who, seed string) string {
 	if babbler, ok := p.babblers[who]; ok {
 		if len(babbler.start.arcs) == 0 {
 			return ""
 		}
-		words := []string{}
-		cur := babbler.start
+		words := []string{seed}
+		var cur *node
+		if cur, ok = babbler.lookup[seed]; !ok {
+			return fmt.Sprintf("%s hasn't used the word '%s'", who, seed)
+		}
 		for cur != babbler.end {
 			which := rand.Intn(cur.wordFrequency)
 			sum := 0
@@ -279,7 +299,7 @@ func (p *BabblerPlugin) babble(who string) string {
 			}
 		}
 
-		return strings.Join(words, " ")
+		return strings.TrimSpace(strings.Join(words, " "))
 	}
 
 	return fmt.Sprintf("could not find babbler: %s", who)
diff --git a/plugins/babbler/babbler_test.go b/plugins/babbler/babbler_test.go
index 5ebc9df..be7aa25 100644
--- a/plugins/babbler/babbler_test.go
+++ b/plugins/babbler/babbler_test.go
@@ -25,6 +25,28 @@ func makeMessage(payload string) msg.Message {
 	}
 }
 
+func TestBabblerNoBabbler(t *testing.T) {
+	mb := bot.NewMockBot()
+	c := New(mb)
+	c.config.Babbler.DefaultUsers = []string{"seabass"}
+	assert.NotNil(t, c)
+	res := c.Message(makeMessage("!seabass2 says"))
+	assert.Len(t, mb.Messages, 1)
+	assert.True(t, res)
+	assert.Contains(t, mb.Messages[0], "could not find babbler: seabass2")
+}
+
+func TestBabblerNothingSaid(t *testing.T) {
+	mb := bot.NewMockBot()
+	c := New(mb)
+	c.config.Babbler.DefaultUsers = []string{"seabass"}
+	assert.NotNil(t, c)
+	res := c.Message(makeMessage("!seabass says"))
+	assert.Len(t, mb.Messages, 1)
+	assert.True(t, res)
+	assert.Contains(t, mb.Messages[0], "Ze ain't said nothin")
+}
+
 func TestBabbler(t *testing.T) {
 	mb := bot.NewMockBot()
 	c := New(mb)
@@ -45,6 +67,44 @@ func TestBabbler(t *testing.T) {
 	assert.Contains(t, mb.Messages[0], "message")
 }
 
+func TestBabblerSeed(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 long"))
+	assert.Len(t, mb.Messages, 1)
+	assert.True(t, res)
+	assert.Contains(t, mb.Messages[0], "long message")
+}
+
+func TestBabblerBadSeed(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 long message"))
+	assert.Len(t, mb.Messages, 1)
+	assert.True(t, res)
+	assert.Contains(t, mb.Messages[0], "try seabass says [seed-token]")
+}
+
 func TestBabblerBatch(t *testing.T) {
 	mb := bot.NewMockBot()
 	c := New(mb)