From 58a0e1719d035d3c2120d7c767595db68b51caf1 Mon Sep 17 00:00:00 2001
From: Dong Chen <dongluo.chen@docker.com>
Date: Wed, 6 Jan 2016 10:33:51 -0800
Subject: [PATCH] Update failureCount scenario and test cases.

Signed-off-by: Dong Chen <dongluo.chen@docker.com>
---
 cluster/engine.go      | 15 +++++++++------
 cluster/engine_test.go | 22 +++++++++++++++++++---
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/cluster/engine.go b/cluster/engine.go
index 0a6af8f47f..c860acfac7 100644
--- a/cluster/engine.go
+++ b/cluster/engine.go
@@ -208,15 +208,11 @@ func (e *Engine) IsHealthy() bool {
 	return e.state == stateHealthy
 }
 
-// setState sets engine healthy state
+// setState sets engine state
 func (e *Engine) setState(state engineState) {
 	e.Lock()
 	defer e.Unlock()
 	e.state = state
-	// if engine is healthy, clear failureCount
-	if state == stateHealthy {
-		e.failureCount = 0
-	}
 }
 
 // TimeToValidate returns true if a pending node is up for validation
@@ -292,10 +288,17 @@ func (e *Engine) UpdatedAt() time.Time {
 	return e.updatedAt
 }
 
+func (e *Engine) resetFailureCount() {
+	e.Lock()
+	defer e.Unlock()
+	e.failureCount = 0
+}
+
 // CheckConnectionErr checks error from client response and adjusts engine healthy indicators
 func (e *Engine) CheckConnectionErr(err error) {
 	if err == nil {
 		e.setErrMsg("")
+		e.resetFailureCount()
 		// If current state is unhealthy, change it to healthy
 		if e.state == stateUnhealthy {
 			log.WithFields(log.Fields{"name": e.Name, "id": e.ID}).Infof("Engine came back to life after %d retries. Hooray!", e.failureCount)
@@ -322,7 +325,7 @@ func (e *Engine) CheckConnectionErr(err error) {
 		e.incFailureCount()
 		return
 	}
-	// other errors may be ambiguous. let refresh loop decide healthy or not.
+	// other errors may be ambiguous.
 }
 
 // Gather engine specs (CPU, memory, constraints, ...).
diff --git a/cluster/engine_test.go b/cluster/engine_test.go
index 2757242513..0cb8fb7762 100644
--- a/cluster/engine_test.go
+++ b/cluster/engine_test.go
@@ -43,11 +43,8 @@ func TestSetEngineState(t *testing.T) {
 	assert.True(t, engine.state == statePending)
 	engine.setState(stateUnhealthy)
 	assert.True(t, engine.state == stateUnhealthy)
-	engine.incFailureCount()
-	assert.True(t, engine.failureCount == 1)
 	engine.setState(stateHealthy)
 	assert.True(t, engine.state == stateHealthy)
-	assert.True(t, engine.failureCount == 0)
 }
 
 func TestErrMsg(t *testing.T) {
@@ -58,6 +55,22 @@ func TestErrMsg(t *testing.T) {
 	assert.True(t, engine.ErrMsg() == message)
 }
 
+func TestCheckConnectionErr(t *testing.T) {
+	engine := NewEngine("test", 0, engOpts)
+	engine.setState(stateHealthy)
+	assert.True(t, engine.failureCount == 0)
+	err := dockerclient.ErrConnectionRefused
+	engine.CheckConnectionErr(err)
+	assert.True(t, len(engine.ErrMsg()) > 0)
+	assert.True(t, engine.failureCount == 1)
+	engine.CheckConnectionErr(err)
+	assert.True(t, engine.failureCount == 2)
+	err = nil
+	engine.CheckConnectionErr(err)
+	assert.True(t, engine.failureCount == 0)
+	assert.True(t, len(engine.ErrMsg()) == 0)
+}
+
 func TestEngineFailureCount(t *testing.T) {
 	engine := NewEngine("test", 0, engOpts)
 	engine.setState(stateHealthy)
@@ -66,6 +79,9 @@ func TestEngineFailureCount(t *testing.T) {
 		engine.incFailureCount()
 	}
 	assert.False(t, engine.IsHealthy())
+	assert.True(t, engine.failureCount == engine.opts.FailureRetry)
+	engine.resetFailureCount()
+	assert.True(t, engine.failureCount == 0)
 }
 
 func TestEngineConnectionFailure(t *testing.T) {