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..33d37ec --- /dev/null +++ b/plugins/stats/stats.go @@ -0,0 +1,216 @@ +// © 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 ( + "bytes" + "encoding/gob" + "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 { + bucket string + key string + val value +} + +// 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) { + var b bytes.Buffer + enc := gob.NewEncoder(&b) + err := enc.Encode(v) + return b.Bytes(), err +} + +func valueFromBytes(b []byte) (value, error) { + var v value + buf := bytes.NewReader(b) + dec := gob.NewDecoder(buf) + err := dec.Decode(&v) + return v, err +} + +type stats []stat + +func mkStat(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{ + 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 +} + +func statFromDB(path, 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() + + b, err := tx.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 + } + + return mkStat(buk, k, v) +} + +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 { + b, err := tx.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{"user", message.User.Name, 1}} +} + +func (p *StatsPlugin) mkHourStat(message msg.Message) stats { + hr := time.Now().Hour() + return stats{stat{"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{"sighting", name, 1}) + } + } + return stats +} + +func (p *StatsPlugin) mkChannelStat(message msg.Message) stats { + return stats{stat{"channel", message.Channel, 1}} +} diff --git a/plugins/stats/stats_test.go b/plugins/stats/stats_test.go new file mode 100644 index 0000000..5d801c5 --- /dev/null +++ b/plugins/stats/stats_test.go @@ -0,0 +1,177 @@ +package stats + +import ( + "encoding/json" + "fmt" + "os" + "strings" + "testing" +) + +func TestJSON(t *testing.T) { + expected := 5 + b, err := json.Marshal(expected) + if err != nil { + t.Error(err) + } + t.Logf("%+v", expected) + t.Log(string(b)) +} + +func TestValueConversion(t *testing.T) { + expected := value(5) + + b, err := expected.Bytes() + if err != nil { + t.Error(err) + return + } + + t.Log(string(b)) + + actual, err := valueFromBytes(b) + if err != nil { + t.Error(err) + return + } + + if fmt.Sprintf("%+v", actual) != fmt.Sprintf("%+v", expected) { + t.Errorf("Did not get equivalent objects: %+v != %+v", actual, expected) + } +} + +func rmDB(t *testing.T, dbPath string) { + err := os.Remove(dbPath) + if err != nil && !strings.Contains(err.Error(), "no such file or directory") { + t.Fatal(err) + } +} + +func TestWithDB(t *testing.T) { + dbPath := "test.db" + rmDB(t, dbPath) + + t.Run("TestDBReadWrite", func(t *testing.T) { + bucket := "testBucket" + key := "testKey" + + expected := stats{stat{ + bucket, + key, + 1, + }} + + err := expected.toDB(dbPath) + if err != nil { + t.Fatalf("Error writing to DB: %s", err) + } + + actual, err := statFromDB(dbPath, bucket, key) + if err != nil { + t.Fatalf("Error reading DB: %s", err) + } + + if actual != expected[0] { + t.Fatalf("%+v != %+v", actual, expected) + } + + }) + + rmDB(t, dbPath) + + t.Run("TestDBAddStatInLoop", func(t *testing.T) { + bucket := "testBucket" + key := "testKey" + expected := value(25) + + os.Remove(dbPath) + defer func() { + err := os.Remove(dbPath) + if err != nil { + t.Fatal(err) + } + }() + + statPack := stats{stat{ + bucket, + key, + 5, + }} + + for i := 0; i < 5; i++ { + err := statPack.toDB(dbPath) + if err != nil { + t.Fatalf("Error writing to DB: %s", err) + } + } + + actual, err := statFromDB(dbPath, bucket, key) + if err != nil { + t.Fatalf("Error reading DB: %s", err) + } + + if actual.val != expected { + t.Fatalf("%+v != %+v", actual.val, expected) + } + }) + + rmDB(t, dbPath) + + t.Run("TestDBAddStats", func(t *testing.T) { + bucket := "testBucket" + key := "testKey" + expected := value(5) + + os.Remove(dbPath) + defer func() { + err := os.Remove(dbPath) + if err != nil { + t.Fatal(err) + } + }() + + statPack := stats{ + stat{ + bucket, + key, + 1, + }, + stat{ + bucket, + key, + 1, + }, + stat{ + bucket, + key, + 1, + }, + stat{ + bucket, + key, + 1, + }, + stat{ + bucket, + key, + 1, + }, + } + + err := statPack.toDB(dbPath) + if err != nil { + t.Fatalf("Error writing to DB: %s", err) + } + + actual, err := statFromDB(dbPath, bucket, key) + if err != nil { + t.Fatalf("Error reading DB: %s", err) + } + + if actual.val != expected { + t.Fatalf("%+v != %+v", actual.val, expected) + } + }) + + rmDB(t, dbPath) +}