diff --git a/storage/mongo.go b/storage/mongo.go index ed2fec6..2bf0bc5 100644 --- a/storage/mongo.go +++ b/storage/mongo.go @@ -2,7 +2,6 @@ package storage import ( "context" - "local/cheqbooq/config" "log" "os/exec" "time" @@ -13,6 +12,7 @@ import ( type Mongo struct { client *mongo.Client + ns string } func init() { @@ -44,7 +44,7 @@ func init() { }() } -func NewMongo(addr string) (*Mongo, error) { +func NewMongo(ns, addr string) (*Mongo, error) { ctx, can := context.WithTimeout(context.Background(), time.Second*5) defer can() @@ -58,6 +58,7 @@ func NewMongo(addr string) (*Mongo, error) { return &Mongo{ client: client, + ns: ns, }, client.Ping(ctx, nil) } @@ -66,15 +67,15 @@ func (m *Mongo) Close() error { } func (m *Mongo) Account() *mongo.Collection { - return m.client.Database(config.StoreNS).Collection("account") + return m.client.Database(m.ns).Collection("account") } func (m *Mongo) Balance() *mongo.Collection { - return m.client.Database(config.StoreNS).Collection("balance") + return m.client.Database(m.ns).Collection("balance") } func (m *Mongo) Transaction() *mongo.Collection { - return m.client.Database(config.StoreNS).Collection("transaction") + return m.client.Database(m.ns).Collection("transaction") } func (m *Mongo) Find(c *mongo.Collection, where interface{}, next func() interface{}) error { diff --git a/storage/mongo_test.go b/storage/mongo_test.go index 76fff3c..489a85f 100644 --- a/storage/mongo_test.go +++ b/storage/mongo_test.go @@ -1,10 +1,97 @@ package storage -import "testing" +import ( + "context" + "crypto/rand" + "encoding/base64" + "testing" +) -func TestMongoNew(t *testing.T) { - _, err := NewMongo("mongodb://localhost:27017") +func testMongoNew(t *testing.T) *Mongo { + b := make([]byte, 5) + rand.Read(b) + m, err := NewMongo(base64.URLEncoding.EncodeToString(b), "mongodb://localhost:27017") if err != nil { t.Fatal(err) } + return m +} + +func TestMongoNew(t *testing.T) { + if err := testMongoNew(t).Close(); err != nil { + t.Fatal(err) + } +} + +func TestMongoFind(t *testing.T) { + m := testMongoNew(t) + defer m.Close() + + c := m.Account() + defer c.Drop(context.TODO()) + if _, err := c.InsertMany(context.TODO(), []interface{}{ + map[string]interface{}{"_id": "1", "a": "b"}, + map[string]interface{}{"_id": "2", "a": "b", "c": "d"}, + map[string]interface{}{"_id": "3", "c": "d"}, + }); err != nil { + t.Fatal(err) + } + + cnt := 0 + inc := func() interface{} { + cnt += 1 + var v interface{} + return &v + } + if err := m.Find(c, map[string]interface{}{"a": "b"}, inc); err != nil { + t.Fatal(err) + } else if cnt != 2 { + t.Fatal(cnt) + } + + cnt = 0 + if err := m.Find(c, map[string]interface{}{"_id": "1"}, inc); err != nil { + t.Fatal(err) + } else if cnt != 1 { + t.Fatal(cnt) + } + + cnt = 0 + if err := m.Find(c, map[string]interface{}{}, inc); err != nil { + t.Fatal(err) + } else if cnt != 3 { + t.Fatal(cnt) + } +} + +func TestMongoUpsert(t *testing.T) { + m := testMongoNew(t) + defer m.Close() + + c := m.Account() + defer c.Drop(context.TODO()) + if _, err := c.InsertMany(context.TODO(), []interface{}{ + map[string]interface{}{"_id": "1", "a": "b"}, + map[string]interface{}{"_id": "2", "a": "b", "c": "d"}, + map[string]interface{}{"_id": "3", "c": "d"}, + }); err != nil { + t.Fatal(err) + } + + if n, err := c.EstimatedDocumentCount(context.TODO()); err != nil { + t.Fatal(err) + } else if n != 3 { + t.Fatal(n) + } + + if err := m.Upsert(c, map[string]interface{}{"_id": "1"}, map[string]interface{}{"$set": map[string]interface{}{"c": "d"}}); err != nil { + t.Fatal(err) + } + + mapped := map[string]string{} + if err := c.FindOne(context.TODO(), map[string]interface{}{"_id": "1"}).Decode(&mapped); err != nil { + t.Fatal(err) + } else if mapped["c"] != "d" { + t.Fatal(mapped) + } } diff --git a/storage/storage.go b/storage/storage.go index 8fe0fae..752eb55 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -11,7 +11,7 @@ type Storage struct { } func New() (*Storage, error) { - mongo, err := NewMongo(fmt.Sprintf("mongodb://%s", strings.TrimPrefix(config.StoreAddr, "mongodb://"))) + mongo, err := NewMongo(config.StoreNS, fmt.Sprintf("mongodb://%s", strings.TrimPrefix(config.StoreAddr, "mongodb://"))) return &Storage{ mongo: mongo, }, err