Merge pull request #56 from velour/stats

Stats plugin

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

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

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}}
+}

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)
+}