mirror of https://github.com/velour/catbase.git
GitHub
commit
1bb3f17dc1
|
|
@ -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
|
||||
|
|
|
|||
5
main.go
5
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))
|
||||
|
|
|
|||
|
|
@ -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}}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
Loading…
Reference in New Issue