2 files changed, 77 insertions, 2 deletions

diff --git a/workers/worker_list_test.go b/workers/worker_list_test.go
index 0ee64a1..93d29a7 100644
--- a/workers/worker_list_test.go
+++ b/workers/worker_list_test.go
@@ -29,6 +29,7 @@ import (
 
 	gomock "github.com/golang/mock/gomock"
 	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/ginkgo/extensions/table"
 	. "github.com/onsi/gomega"
 	"github.com/satori/go.uuid"
 )
@@ -850,6 +851,72 @@ var _ = Describe("WorkerList", func() {
 				})
 			})
 		})
+		Describe("setState with changeListener", func() {
+			var ctrl *gomock.Controller
+			var wc *MockWorkerChange
+
+			set := func(state WorkerState) {
+				wl.mutex.Lock()
+				wl.workers[worker].State = state
+				wl.mutex.Unlock()
+			}
+			check := func(state WorkerState) {
+				wl.mutex.RLock()
+				Expect(wl.workers[worker].State).To(Equal(state))
+				wl.mutex.RUnlock()
+			}
+			BeforeEach(func() {
+				ctrl = gomock.NewController(GinkgoT())
+				wc = NewMockWorkerChange(ctrl)
+				wl.SetChangeListener(wc)
+				Expect(wl.changeListener).To(Equal(wc))
+			})
+			AfterEach(func() {
+				ctrl.Finish()
+			})
+			DescribeTable("Should change state without calling changeListener",
+				func(from, to WorkerState) {
+					set(from)
+					err := wl.setState(worker, to)
+					Expect(err).NotTo(HaveOccurred())
+					check(to)
+				},
+				Entry("MAINTENANCE->MAINTENANCE", MAINTENANCE, MAINTENANCE),
+				Entry("MAINTENANCE->RUN", MAINTENANCE, RUN),
+				Entry("MAINTENANCE->FAIL", MAINTENANCE, FAIL),
+				Entry("IDLE->MAINTENANCE", IDLE, MAINTENANCE),
+				Entry("IDLE->RUN", IDLE, RUN),
+				Entry("IDLE->FAIL", IDLE, FAIL),
+				Entry("FAIL->MAINTENANCE", FAIL, MAINTENANCE),
+				Entry("FAIL->RUN", FAIL, RUN),
+				Entry("FAIL->FAIL", FAIL, FAIL),
+			)
+			DescribeTable("Should change state and call OnWorkerIdle",
+				func(from, to WorkerState) {
+					set(from)
+					wc.EXPECT().OnWorkerIdle(worker)
+					err := wl.setState(worker, to)
+					Expect(err).NotTo(HaveOccurred())
+					check(to)
+				},
+				Entry("MAINTENANCE->IDLE", MAINTENANCE, IDLE),
+				Entry("IDLE->IDLE", IDLE, IDLE),
+				Entry("RUN->IDLE", RUN, IDLE),
+				Entry("FAIL->IDLE", FAIL, IDLE),
+			)
+			DescribeTable("Should change state and call OnWorkerFail",
+				func(from, to WorkerState) {
+					set(from)
+					wc.EXPECT().OnWorkerFail(worker)
+					err := wl.setState(worker, to)
+					Expect(err).NotTo(HaveOccurred())
+					check(to)
+				},
+				Entry("RUN->MAINTENANCE", RUN, MAINTENANCE),
+				Entry("RUN->RUN", RUN, RUN),
+				Entry("RUN->FAIL", RUN, FAIL),
+			)
+		})
 	})
 	Describe("TakeBestMatchingWorker", func() {
 		addWorker := func(groups Groups, caps Capabilities) *mapWorker {
diff --git a/workers/workers.go b/workers/workers.go
index b6df880..4771226 100644
--- a/workers/workers.go
+++ b/workers/workers.go
@@ -111,8 +111,7 @@ func (wl *WorkerList) SetFail(uuid WorkerUUID, reason string) error {
 	if worker.State == MAINTENANCE {
 		return ErrInMaintenance
 	}
-	worker.State = FAIL
-	return nil
+	return wl.setState(uuid, FAIL)
 }
 
 // SetState is an implementation of SetState from Workers interface.
@@ -398,6 +397,15 @@ func (wl *WorkerList) setState(worker WorkerUUID, state WorkerState) error {
 	if !ok {
 		return ErrWorkerNotFound
 	}
+	if wl.changeListener != nil {
+		if state == IDLE {
+			wl.changeListener.OnWorkerIdle(worker)
+		} else {
+			if w.State == RUN {
+				wl.changeListener.OnWorkerFail(worker)
+			}
+		}
+	}
 	w.State = state
 	return nil
 }