diff --git a/config/config.go b/config/config.go index eafa002..7a0553d 100644 --- a/config/config.go +++ b/config/config.go @@ -75,6 +75,10 @@ type Config struct { Reminder struct { MaxBatchAdd int } + Stats struct { + DBPath string + Sightings []string + } } // Readconfig loads the config data out of a JSON file located in cfile diff --git a/main.go b/main.go index f542e30..ad3c375 100644 --- a/main.go +++ b/main.go @@ -14,11 +14,11 @@ import ( "github.com/velour/catbase/plugins/beers" "github.com/velour/catbase/plugins/counter" "github.com/velour/catbase/plugins/dice" - "github.com/velour/catbase/plugins/downtime" "github.com/velour/catbase/plugins/fact" "github.com/velour/catbase/plugins/leftpad" "github.com/velour/catbase/plugins/reminder" "github.com/velour/catbase/plugins/rss" + "github.com/velour/catbase/plugins/stats" "github.com/velour/catbase/plugins/talker" "github.com/velour/catbase/plugins/your" "github.com/velour/catbase/plugins/zork" @@ -46,9 +46,10 @@ func main() { // b.AddHandler(plugins.NewTestPlugin(b)) b.AddHandler("admin", admin.New(b)) + b.AddHandler("stats", stats.New(b)) // b.AddHandler("first", plugins.NewFirstPlugin(b)) b.AddHandler("leftpad", leftpad.New(b)) - b.AddHandler("downtime", downtime.New(b)) + // b.AddHandler("downtime", downtime.New(b)) b.AddHandler("talker", talker.New(b)) b.AddHandler("dice", dice.New(b)) b.AddHandler("beers", beers.New(b)) diff --git a/plugins/stats/stats.go b/plugins/stats/stats.go new file mode 100644 index 0000000..58d652c --- /dev/null +++ b/plugins/stats/stats.go @@ -0,0 +1,240 @@ +// © 2016 the CatBase Authors under the WTFPL license. See AUTHORS for the list of authors. + +// Stats contains the plugin that allows the bot record chat statistics +package stats + +import ( + "encoding/json" + "log" + "strings" + "time" + + "github.com/boltdb/bolt" + "github.com/velour/catbase/bot" + "github.com/velour/catbase/bot/msg" + "github.com/velour/catbase/config" +) + +type StatsPlugin struct { + bot bot.Bot + config *config.Config +} + +// New creates a new StatsPlugin with the Plugin interface +func New(bot bot.Bot) *StatsPlugin { + p := StatsPlugin{ + bot: bot, + config: bot.Config(), + } + return &p + +} + +type stat struct { + // date formatted: "2006-01-02" + day string + // category + bucket string + // specific unique individual + key string + val value +} + +func mkDay() string { + return time.Now().Format("2006-01-02") +} + +// The value type is here in the future growth case that we might want to put a +// struct of more interesting information into the DB +type value int + +func (v value) Bytes() ([]byte, error) { + b, err := json.Marshal(v) + return b, err +} + +func valueFromBytes(b []byte) (value, error) { + var v value + err := json.Unmarshal(b, &v) + return v, err +} + +type stats []stat + +// mkStat converts raw data to a stat struct +// Expected a string representation of the date formatted: "2006-01-02" +func mkStat(day string, bucket, key, val []byte) (stat, error) { + v, err := valueFromBytes(val) + if err != nil { + log.Printf("mkStat: error getting value from bytes: %s", err) + return stat{}, err + } + return stat{ + day: day, + bucket: string(bucket), + key: string(key), + val: v, + }, nil +} + +// Another future-proofing function I shouldn't have written +func (v value) add(other value) value { + return v + other +} + +// statFromDB takes a location specification and returns the data at that path +// Expected a string representation of the date formatted: "2006-01-02" +func statFromDB(path, day, bucket, key string) (stat, error) { + db, err := bolt.Open(path, 0600, &bolt.Options{ + Timeout: 1 * time.Second, + }) + buk := []byte(bucket) + k := []byte(key) + if err != nil { + log.Printf("Couldn't open BoltDB for stats: %s", err) + return stat{}, err + } + defer db.Close() + + tx, err := db.Begin(true) + if err != nil { + log.Println("statFromDB: Error beginning the Tx") + return stat{}, err + } + defer tx.Rollback() + + d, err := tx.CreateBucketIfNotExists([]byte(day)) + if err != nil { + log.Println("statFromDB: Error creating the bucket") + return stat{}, err + } + b, err := d.CreateBucketIfNotExists(buk) + if err != nil { + log.Println("statFromDB: Error creating the bucket") + return stat{}, err + } + + v := b.Get(k) + + if err := tx.Commit(); err != nil { + log.Println("statFromDB: Error commiting the Tx") + return stat{}, err + } + + if v == nil { + return stat{day, bucket, key, 0}, nil + } + + return mkStat(day, buk, k, v) +} + +// toDB takes a stat and records it, adding to the value in the DB if necessary +func (s stats) toDB(path string) error { + db, err := bolt.Open(path, 0600, &bolt.Options{Timeout: 1 * time.Second}) + if err != nil { + log.Printf("Couldn't open BoltDB for stats: %s", err) + return err + } + defer db.Close() + + for _, stat := range s { + err = db.Update(func(tx *bolt.Tx) error { + d, err := tx.CreateBucketIfNotExists([]byte(stat.day)) + if err != nil { + log.Println("toDB: Error creating bucket") + return err + } + b, err := d.CreateBucketIfNotExists([]byte(stat.bucket)) + if err != nil { + log.Println("toDB: Error creating bucket") + return err + } + + valueInDB := b.Get([]byte(stat.key)) + if valueInDB != nil { + val, err := valueFromBytes(valueInDB) + if err != nil { + log.Println("toDB: Error getting value from bytes") + return err + } + stat.val = stat.val.add(val) + } + + v, err := stat.val.Bytes() + if err != nil { + return err + } + err = b.Put([]byte(stat.key), v) + return err + }) + if err != nil { + return err + } + } + return nil +} + +func (p *StatsPlugin) record(message msg.Message) { + statGenerators := []func(msg.Message) stats{ + p.mkUserStat, + p.mkHourStat, + p.mkChannelStat, + p.mkSightingStat, + } + + allStats := stats{} + + for _, mk := range statGenerators { + stats := mk(message) + if stats != nil { + allStats = append(allStats, stats...) + } + } + + allStats.toDB(p.bot.Config().Stats.DBPath) +} + +func (p *StatsPlugin) Message(message msg.Message) bool { + p.record(message) + return false +} + +func (p *StatsPlugin) Event(e string, message msg.Message) bool { + p.record(message) + return false +} + +func (p *StatsPlugin) BotMessage(message msg.Message) bool { + p.record(message) + return false +} + +func (p *StatsPlugin) Help(e string, m []string) { +} + +func (p *StatsPlugin) RegisterWeb() *string { + return nil +} + +func (p *StatsPlugin) mkUserStat(message msg.Message) stats { + return stats{stat{mkDay(), "user", message.User.Name, 1}} +} + +func (p *StatsPlugin) mkHourStat(message msg.Message) stats { + hr := time.Now().Hour() + return stats{stat{mkDay(), "user", string(hr), 1}} +} + +func (p *StatsPlugin) mkSightingStat(message msg.Message) stats { + stats := stats{} + for _, name := range p.bot.Config().Stats.Sightings { + if strings.Contains(message.Body, name+" sighting") { + stats = append(stats, stat{mkDay(), "sighting", name, 1}) + } + } + return stats +} + +func (p *StatsPlugin) mkChannelStat(message msg.Message) stats { + return stats{stat{mkDay(), "channel", message.Channel, 1}} +} diff --git a/plugins/stats/stats_test.go b/plugins/stats/stats_test.go new file mode 100644 index 0000000..d5597f4 --- /dev/null +++ b/plugins/stats/stats_test.go @@ -0,0 +1,290 @@ +package stats + +import ( + "encoding/json" + "os" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/velour/catbase/bot" + "github.com/velour/catbase/bot/msg" + "github.com/velour/catbase/bot/user" +) + +var dbPath = "test.db" + +func TestJSON(t *testing.T) { + expected := 5 + b, err := json.Marshal(expected) + assert.Nil(t, err) + t.Logf("%+v", expected) + t.Log(string(b)) +} + +func TestValueConversion(t *testing.T) { + expected := value(5) + + b, err := expected.Bytes() + assert.Nil(t, err) + + t.Log(string(b)) + + actual, err := valueFromBytes(b) + assert.Nil(t, err) + + assert.Equal(t, actual, expected) +} + +func rmDB(t *testing.T) { + err := os.Remove(dbPath) + if err != nil && !strings.Contains(err.Error(), "no such file or directory") { + t.Fatal(err) + } +} + +func TestWithDB(t *testing.T) { + rmDB(t) + + t.Run("TestDBReadWrite", func(t *testing.T) { + day := mkDay() + bucket := "testBucket" + key := "testKey" + + expected := stats{stat{ + day, + bucket, + key, + 1, + }} + + err := expected.toDB(dbPath) + assert.Nil(t, err) + + actual, err := statFromDB(dbPath, day, bucket, key) + assert.Nil(t, err) + + assert.Equal(t, actual, expected[0]) + + }) + + rmDB(t) + + t.Run("TestDBAddStatInLoop", func(t *testing.T) { + day := mkDay() + bucket := "testBucket" + key := "testKey" + expected := value(25) + + statPack := stats{stat{ + day, + bucket, + key, + 5, + }} + + for i := 0; i < 5; i++ { + err := statPack.toDB(dbPath) + assert.Nil(t, err) + } + + actual, err := statFromDB(dbPath, day, bucket, key) + assert.Nil(t, err) + + assert.Equal(t, actual.val, expected) + }) + + rmDB(t) + + t.Run("TestDBAddStats", func(t *testing.T) { + day := mkDay() + bucket := "testBucket" + key := "testKey" + expected := value(5) + + statPack := stats{} + for i := 0; i < 5; i++ { + statPack = append(statPack, stat{ + day, + bucket, + key, + 1, + }) + } + + err := statPack.toDB(dbPath) + assert.Nil(t, err) + + actual, err := statFromDB(dbPath, day, bucket, key) + assert.Nil(t, err) + + assert.Equal(t, actual.val, expected) + }) + + rmDB(t) +} + +func makeMessage(payload string) msg.Message { + isCmd := strings.HasPrefix(payload, "!") + if isCmd { + payload = payload[1:] + } + return msg.Message{ + User: &user.User{Name: "tester"}, + Channel: "test", + Body: payload, + Command: isCmd, + } +} + +func testUserCounter(t *testing.T, count int) { + day := mkDay() + expected := value(count) + mb := bot.NewMockBot() + mb.Cfg.Stats.DBPath = dbPath + s := New(mb) + assert.NotNil(t, s) + + for i := 0; i < count; i++ { + s.Message(makeMessage("test")) + } + + _, err := os.Stat(dbPath) + assert.Nil(t, err) + + stat, err := statFromDB(mb.Config().Stats.DBPath, day, "user", "tester") + assert.Nil(t, err) + actual := stat.val + assert.Equal(t, actual, expected) +} + +func TestMessages(t *testing.T) { + _, err := os.Stat(dbPath) + assert.NotNil(t, err) + + t.Run("TestOneUserCounter", func(t *testing.T) { + day := mkDay() + count := 5 + expected := value(count) + mb := bot.NewMockBot() + mb.Cfg.Stats.DBPath = dbPath + s := New(mb) + assert.NotNil(t, s) + + for i := 0; i < count; i++ { + s.Message(makeMessage("test")) + } + + _, err := os.Stat(dbPath) + assert.Nil(t, err) + + stat, err := statFromDB(mb.Config().Stats.DBPath, day, "user", "tester") + assert.Nil(t, err) + actual := stat.val + assert.Equal(t, actual, expected) + }) + + rmDB(t) + + t.Run("TestTenUserCounter", func(t *testing.T) { + day := mkDay() + count := 5 + expected := value(count) + mb := bot.NewMockBot() + mb.Cfg.Stats.DBPath = dbPath + s := New(mb) + assert.NotNil(t, s) + + for i := 0; i < count; i++ { + s.Message(makeMessage("test")) + } + + _, err := os.Stat(dbPath) + assert.Nil(t, err) + + stat, err := statFromDB(mb.Config().Stats.DBPath, day, "user", "tester") + assert.Nil(t, err) + actual := stat.val + assert.Equal(t, actual, expected) + }) + + rmDB(t) + + t.Run("TestChannelCounter", func(t *testing.T) { + day := mkDay() + count := 5 + expected := value(count) + mb := bot.NewMockBot() + mb.Cfg.Stats.DBPath = dbPath + s := New(mb) + assert.NotNil(t, s) + + for i := 0; i < count; i++ { + s.Message(makeMessage("test")) + } + + _, err := os.Stat(dbPath) + assert.Nil(t, err) + + stat, err := statFromDB(mb.Config().Stats.DBPath, day, "channel", "test") + assert.Nil(t, err) + actual := stat.val + assert.Equal(t, actual, expected) + }) + + rmDB(t) + + t.Run("TestSightingCounter", func(t *testing.T) { + day := mkDay() + count := 5 + expected := value(count) + mb := bot.NewMockBot() + + mb.Cfg.Stats.DBPath = dbPath + mb.Cfg.Stats.Sightings = []string{"user", "nobody"} + + s := New(mb) + assert.NotNil(t, s) + + for i := 0; i < count; i++ { + s.Message(makeMessage("user sighting")) + } + + _, err := os.Stat(dbPath) + assert.Nil(t, err) + + stat, err := statFromDB(mb.Config().Stats.DBPath, day, "sighting", "user") + assert.Nil(t, err) + actual := stat.val + assert.Equal(t, actual, expected) + }) + + rmDB(t) + + t.Run("TestSightingCounterNoResults", func(t *testing.T) { + day := mkDay() + count := 5 + expected := value(0) + mb := bot.NewMockBot() + + mb.Cfg.Stats.DBPath = dbPath + mb.Cfg.Stats.Sightings = []string{} + + s := New(mb) + assert.NotNil(t, s) + + for i := 0; i < count; i++ { + s.Message(makeMessage("user sighting")) + } + + _, err := os.Stat(dbPath) + assert.Nil(t, err) + + stat, err := statFromDB(mb.Config().Stats.DBPath, day, "sighting", "user") + assert.Nil(t, err) + actual := stat.val + assert.Equal(t, actual, expected) + }) + + rmDB(t) +}