// PutTask stores a requested task in the database func (db *TaskDB) PutTask(task *eremetic.Task) error { return db.conn.Update(func(tx *bolt.Tx) error { b, err := tx.CreateBucketIfNotExists([]byte("tasks")) if err != nil { return err } encoded, err := eremetic.Encode(task) if err != nil { logrus.WithError(err).Error("Unable to encode task to byte-array.") return err } return b.Put([]byte(task.ID), encoded) }) }
func TestZKDatabase(t *testing.T) { var ( db *TaskDB object *mockConnection connector *mockConnector ) zkPath := "zk://localhost:1234/testdb" setup := func() { db = &TaskDB{ conn: new(mockConnection), path: "/testdb", } object = db.conn.(*mockConnection) connector = new(mockConnector) } teardown := func() { db = nil object = nil connector = nil } status := []eremetic.Status{ eremetic.Status{ Status: eremetic.TaskRunning, Time: time.Now().Unix(), }, } var maskedEnv = make(map[string]string) maskedEnv["foo"] = "bar" task := &eremetic.Task{ ID: "1234", MaskedEnvironment: maskedEnv, Status: status, } taskBytes, err := eremetic.Encode(task) if err != nil { t.Fail() } Convey("Creating", t, func() { Convey("Errors", func() { Convey("Missing path", func() { setup() defer teardown() _, err := newCustomTaskDB(connector, "") So(err, ShouldNotBeNil) So(err.Error(), ShouldEqual, "Missing ZK path") }) Convey("Unable to connect", func() { setup() defer teardown() connector.On("Connect", mock.AnythingOfType("string")).Return(nil, errors.New("Unable to connect")) _, err := newCustomTaskDB(connector, zkPath) So(err, ShouldNotBeNil) So(connector.AssertCalled(t, "Connect", "localhost:1234"), ShouldBeTrue) }) Convey("Unable to verify existance", func() { setup() defer teardown() connector.On("Connect", mock.AnythingOfType("string")).Return(object, nil) object.On("Exists", mock.AnythingOfType("string")).Return(false, &zk.Stat{}, errors.New("Bad Connection")) _, err := newCustomTaskDB(connector, zkPath) So(err, ShouldNotBeNil) So(err.Error(), ShouldEqual, "Bad Connection") So(connector.AssertCalled(t, "Connect", "localhost:1234"), ShouldBeTrue) So(object.AssertCalled(t, "Exists", "/testdb"), ShouldBeTrue) }) Convey("Fail to create if not exists", func() { setup() defer teardown() connector.On("Connect", mock.AnythingOfType("string")).Return(object, nil) object.On("Exists", mock.AnythingOfType("string")).Return(false, &zk.Stat{}, nil) object.On("Create", mock.AnythingOfType("string"), mock.Anything, mock.AnythingOfType("int32"), mock.Anything).Return("", errors.New("Unable to create node")) _, err := newCustomTaskDB(connector, zkPath) So(err, ShouldNotBeNil) So(err.Error(), ShouldEqual, "Unable to create node") So(connector.AssertCalled(t, "Connect", "localhost:1234"), ShouldBeTrue) So(object.AssertCalled(t, "Exists", "/testdb"), ShouldBeTrue) So(object.AssertCalled(t, "Create", "/testdb", mock.Anything, mock.AnythingOfType("int32"), mock.Anything), ShouldBeTrue) }) }) Convey("Success", func() { setup() defer teardown() connector.On("Connect", mock.AnythingOfType("string")).Return(object, nil) object.On("Exists", mock.AnythingOfType("string")).Return(true, &zk.Stat{}, nil) db, err := newCustomTaskDB(connector, zkPath) So(err, ShouldBeNil) So(connector.AssertCalled(t, "Connect", "localhost:1234"), ShouldBeTrue) So(object.AssertCalled(t, "Exists", "/testdb"), ShouldBeTrue) So(db, ShouldImplement, (*eremetic.TaskDB)(nil)) // Most weirdest syntax ever? }) }) Convey("Clean", t, func() { Convey("Success", func() { setup() defer teardown() object.On("Delete", mock.AnythingOfType("string"), mock.AnythingOfType("int32")).Return(nil) err := db.Clean() So(err, ShouldBeNil) So(object.AssertCalled(t, "Delete", "/testdb/", mock.Anything), ShouldBeTrue) }) }) Convey("Close", t, func() { setup() defer teardown() object.On("Close").Return(nil) db.Close() So(object.AssertCalled(t, "Close"), ShouldBeTrue) }) Convey("AddTask", t, func() { Convey("Exists", func() { setup() defer teardown() object.On("Exists", mock.AnythingOfType("string")).Return(true, &zk.Stat{}, nil) object.On("Set", mock.Anything, mock.Anything, mock.Anything).Return(&zk.Stat{}, nil) err := db.PutTask(task) So(err, ShouldBeNil) So(object.AssertCalled(t, "Exists", "/testdb/1234"), ShouldBeTrue) So(object.AssertCalled(t, "Set", "/testdb/1234", taskBytes, mock.AnythingOfType("int32")), ShouldBeTrue) }) Convey("New", func() { setup() defer teardown() object.On("Exists", mock.AnythingOfType("string")).Return(false, &zk.Stat{}, nil) object.On("Create", mock.AnythingOfType("string"), mock.Anything, mock.AnythingOfType("int32"), mock.Anything).Return("", nil) err := db.PutTask(task) So(err, ShouldBeNil) So(object.AssertCalled(t, "Exists", "/testdb/1234"), ShouldBeTrue) So(object.AssertCalled(t, "Create", "/testdb/1234", taskBytes, mock.AnythingOfType("int32"), mock.Anything), ShouldBeTrue) }) Convey("Errors", func() { Convey("Bad Connection", func() { setup() defer teardown() object.On("Exists", mock.AnythingOfType("string")).Return(false, &zk.Stat{}, errors.New("Bad Connection")) err := db.PutTask(task) So(err.Error(), ShouldEqual, "Bad Connection") }) }) }) Convey("ReadUnmaskedTask", t, func() { Convey("Success", func() { setup() defer teardown() object.On("Get", mock.AnythingOfType("string")).Return(taskBytes, &zk.Stat{}, nil) read, err := db.ReadUnmaskedTask("1234") So(err, ShouldBeNil) So(&read, ShouldResemble, task) So(read.MaskedEnvironment["foo"], ShouldEqual, "bar") So(object.AssertCalled(t, "Get", "/testdb/1234"), ShouldBeTrue) }) Convey("Error", func() { setup() defer teardown() object.On("Get", mock.AnythingOfType("string")).Return([]byte{}, &zk.Stat{}, errors.New("Unable to Read")) _, err := db.ReadUnmaskedTask("1234") So(err.Error(), ShouldEqual, "Unable to Read") So(object.AssertCalled(t, "Get", "/testdb/1234"), ShouldBeTrue) }) }) Convey("ReadTask", t, func() { Convey("Success", func() { setup() defer teardown() object.On("Get", mock.AnythingOfType("string")).Return(taskBytes, &zk.Stat{}, nil) read, err := db.ReadTask("1234") So(err, ShouldBeNil) So(&read, ShouldHaveSameTypeAs, task) So(read.ID, ShouldEqual, task.ID) So(read.MaskedEnvironment["foo"], ShouldEqual, eremetic.Masking) So(object.AssertCalled(t, "Get", "/testdb/1234"), ShouldBeTrue) }) Convey("Error", func() { setup() defer teardown() object.On("Get", mock.AnythingOfType("string")).Return([]byte{}, &zk.Stat{}, errors.New("Unable to Read")) _, err := db.ReadTask("1234") So(err.Error(), ShouldEqual, "Unable to Read") So(object.AssertCalled(t, "Get", "/testdb/1234"), ShouldBeTrue) }) }) Convey("ListNonTerminalTasks", t, func() { Convey("Success", func() { setup() defer teardown() object.On("Children", mock.AnythingOfType("string")).Return([]string{"1234"}, nil, nil) object.On("Get", mock.AnythingOfType("string")).Return(taskBytes, &zk.Stat{}, nil) list, err := db.ListNonTerminalTasks() So(err, ShouldBeNil) So(list, ShouldHaveLength, 1) So(list[0], ShouldHaveSameTypeAs, task) So(list[0].ID, ShouldEqual, task.ID) So(list[0].MaskedEnvironment["foo"], ShouldEqual, eremetic.Masking) So(object.AssertCalled(t, "Get", "/testdb/1234"), ShouldBeTrue) So(object.AssertCalled(t, "Children", "/testdb"), ShouldBeTrue) }) Convey("Error", func() { setup() defer teardown() object.On("Children", mock.AnythingOfType("string")).Return([]string{"1234"}, nil, nil) object.On("Get", mock.AnythingOfType("string")).Return([]byte{}, &zk.Stat{}, errors.New("Unable to Read")) list, _ := db.ListNonTerminalTasks() So(list, ShouldBeEmpty) So(list, ShouldNotBeNil) So(object.AssertCalled(t, "Get", "/testdb/1234"), ShouldBeTrue) So(object.AssertCalled(t, "Children", "/testdb"), ShouldBeTrue) }) }) Convey("parsePath", t, func() { masters := make(map[string]string) masters["master1.local:1111,master2.local:1111,master3.local:1111"] = "zk://master1.local:1111,master2.local:1111,master3.local:1111/mesos" masters["master1.local:2222,master2.local:2222,master3.local:2222"] = "zk://master1.local:2222,master2.local:2222,master3.local:2222/cluster/mesos" masters["master1.local:3333,master2.local:3333,master3.local:3333"] = "zk://master1.local:3333,master2.local:3333,master3.local:3333/mesos/cluster" masters["123.123.123.123:4444,10.1.15.4:4444,10.1.15.42:4444"] = "zk://123.123.123.123:4444,10.1.15.4:4444,10.1.15.42:4444/" for e, p := range masters { Convey(p, func() { servers, paths, err := parsePath(p) So(err, ShouldBeNil) So(servers, ShouldEqual, e) So(strings.HasSuffix(p, paths), ShouldBeTrue) }) } Convey("Ensures it starts with a / but doesn't end with a /", func() { _, path, err := parsePath("zk://master1.local:1111/pathing/") So(err, ShouldBeNil) So(strings.HasPrefix(path, "/"), ShouldBeTrue) So(strings.HasSuffix(path, "/"), ShouldBeFalse) }) }) }