From e6324ad5a2b276e3dd4d9732303842aa40ee9c92 Mon Sep 17 00:00:00 2001 From: Chris Sexton Date: Tue, 10 May 2016 21:15:52 -0400 Subject: [PATCH] babbler: Add all users in chan, add cfg'd users Also misc small updates to the mock object for easier config usage --- bot/mock.go | 4 +- config/config.go | 3 + example_config.json | 5 ++ plugins/babbler/babbler.go | 98 ++++++++++++++++++++------------- plugins/babbler/babbler_test.go | 2 + slack/slack.go | 11 +++- 6 files changed, 82 insertions(+), 41 deletions(-) diff --git a/bot/mock.go b/bot/mock.go index 1655155..193c5ee 100644 --- a/bot/mock.go +++ b/bot/mock.go @@ -16,11 +16,13 @@ type MockBot struct { mock.Mock db *sqlx.DB + Cfg config.Config + Messages []string Actions []string } -func (mb *MockBot) Config() *config.Config { return &config.Config{} } +func (mb *MockBot) Config() *config.Config { return &mb.Cfg } func (mb *MockBot) DBVersion() int64 { return 1 } func (mb *MockBot) DB() *sqlx.DB { return mb.db } func (mb *MockBot) Who(string) []user.User { return []user.User{} } diff --git a/config/config.go b/config/config.go index df4ee2b..112d279 100644 --- a/config/config.go +++ b/config/config.go @@ -64,6 +64,9 @@ type Config struct { QuoteTime int StartupFact string } + Babbler struct { + DefaultUsers []string + } } // Readconfig loads the config data out of a JSON file located in cfile diff --git a/example_config.json b/example_config.json index 08eb01b..3540c3f 100644 --- a/example_config.json +++ b/example_config.json @@ -58,5 +58,10 @@ "QuoteTime": 1, "StartupFact": "speed test", "MinLen": 5 + }, + "Babbler": { + "DefaultUsers": [ + "seabass" + ] } } diff --git a/plugins/babbler/babbler.go b/plugins/babbler/babbler.go index 144b435..685bd04 100644 --- a/plugins/babbler/babbler.go +++ b/plugins/babbler/babbler.go @@ -8,59 +8,79 @@ import ( "math/rand" "strings" - "github.com/jmoiron/sqlx" "github.com/velour/catbase/bot" "github.com/velour/catbase/bot/msg" + "github.com/velour/catbase/bot/user" + "github.com/velour/catbase/config" ) type BabblerPlugin struct { - Bot bot.Bot - db *sqlx.DB + Bot bot.Bot + db *sqlx.DB + config *config.Config babblers map[string]*babbler } type babbler struct { - start *node - end *node + start *node + end *node lookup map[string]*node } type node struct { wordFrequency int - arcs map[string]*arc + arcs map[string]*arc } type arc struct { transitionFrequency int - next *node + next *node } func New(bot bot.Bot) *BabblerPlugin { plugin := &BabblerPlugin{ - Bot: bot, - db: bot.DB(), + Bot: bot, + db: bot.DB(), + config: bot.Config(), babblers: map[string]*babbler{}, } - // this who string isn't escaped, just sooo, you know. - babbler, err := getMarkovChain(plugin.db, "seabass") - if err == nil { - plugin.babblers["seabass"] = babbler - } else { - plugin.babblers["seabass"] = newBabbler() - } - return plugin } +func (p *BabblerPlugin) makeBabbler(newUser user.User) { + name := newUser.Name + babbler, err := getMarkovChain(p.db, name) + if err == nil { + p.babblers[name] = babbler + } else { + p.babblers[name] = newBabbler() + } +} + +func (p *BabblerPlugin) makeBabblers(newUser user.User) { + users := p.Bot.Who(p.config.MainChannel) + users = append(users, newUser) + for _, name := range p.config.Babbler.DefaultUsers { + users = append(users, user.New(name)) + } + for _, u := range users { + p.makeBabbler(u) + } +} + func (p *BabblerPlugin) Message(message msg.Message) bool { + if len(p.babblers) == 0 { + p.makeBabblers(*message.User) + } else if _, ok := p.babblers[message.User.Name]; !ok { + p.makeBabbler(*message.User) + } + lowercase := strings.ToLower(message.Body) tokens := strings.Fields(lowercase) - if _, ok := p.babblers[message.User.Name]; ok { - addToMarkovChain(p.babblers[message.User.Name], lowercase) - } + addToMarkovChain(p.babblers[message.User.Name], lowercase) if len(tokens) == 4 && strings.Contains(lowercase, "initialize babbler for ") { who := tokens[len(tokens)-1] @@ -71,20 +91,24 @@ func (p *BabblerPlugin) Message(message msg.Message) bool { } else { p.babblers[who] = newBabbler() } + p.Bot.SendMessage(message.Channel, "Okay.") + return true } } if len(tokens) == 2 && tokens[1] == "says" { - if _, ok := p.babblers[tokens[0]]; ok { - p.Bot.SendMessage(message.Channel, p.babble(tokens[0])) - return true + saying := p.babble(tokens[0]) + if saying == "" { + p.Bot.SendMessage(message.Channel, "Ze ain't said nothin'") } + p.Bot.SendMessage(message.Channel, saying) + return true } return false } func (p *BabblerPlugin) Help(channel string, parts []string) { - p.Bot.SendMessage(channel, "seabass says") + p.Bot.SendMessage(channel, "initialize babbler for seabass\n\nseabass says") } func (p *BabblerPlugin) Event(kind string, message msg.Message) bool { @@ -109,7 +133,7 @@ func addToMarkovChain(babble *babbler, phrase string) { if _, ok := babble.lookup[words[i]]; !ok { babble.lookup[words[i]] = &node{ wordFrequency: 1, - arcs: map[string]*arc{}, + arcs: map[string]*arc{}, } } else { babble.lookup[words[i]].wordFrequency++ @@ -119,7 +143,7 @@ func addToMarkovChain(babble *babbler, phrase string) { if _, ok := prev.arcs[words[i]]; !ok { prev.arcs[words[i]] = &arc{ transitionFrequency: 1, - next: babble.lookup[words[i]], + next: babble.lookup[words[i]], } } else { prev.arcs[words[i]].transitionFrequency++ @@ -131,7 +155,7 @@ func addToMarkovChain(babble *babbler, phrase string) { if _, ok := prev.arcs[""]; !ok { prev.arcs[""] = &arc{ transitionFrequency: 1, - next: babble.end, + next: babble.end, } } else { prev.arcs[""].transitionFrequency++ @@ -139,15 +163,15 @@ func addToMarkovChain(babble *babbler, phrase string) { } func newBabbler() *babbler { - return &babbler { - start: &node { - wordFrequency: 0, - arcs: map[string]*arc{}, - }, - end: &node { - wordFrequency: 0, - arcs: map[string]*arc{}, - }, + return &babbler{ + start: &node{ + wordFrequency: 0, + arcs: map[string]*arc{}, + }, + end: &node{ + wordFrequency: 0, + arcs: map[string]*arc{}, + }, lookup: map[string]*node{}, } } @@ -199,4 +223,4 @@ func (p *BabblerPlugin) babble(who string) string { } return fmt.Sprintf("could not find babbler: %s", who) -} \ No newline at end of file +} diff --git a/plugins/babbler/babbler_test.go b/plugins/babbler/babbler_test.go index da93362..aabbd7d 100644 --- a/plugins/babbler/babbler_test.go +++ b/plugins/babbler/babbler_test.go @@ -28,10 +28,12 @@ func makeMessage(payload string) msg.Message { func TestBabbler(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" diff --git a/slack/slack.go b/slack/slack.go index 10adcc1..ead635f 100644 --- a/slack/slack.go +++ b/slack/slack.go @@ -258,9 +258,14 @@ func (s *Slack) Who(id string) []string { u := s.url + "channels.info" resp, err := http.PostForm(u, url.Values{"token": {s.config.Slack.Token}, "channel": {id}}) - if err != nil || resp.StatusCode != 200 { - log.Printf("Error posting user info request: %d %s", - resp.StatusCode, err) + if err != nil { + log.Printf("Error posting user info request: %s", + err) + return []string{} + } + if resp.StatusCode != 200 { + log.Printf("Error posting user info request: %d", + resp.StatusCode) return []string{} } defer resp.Body.Close()