polish tests

change_test.go

diff --git a/change_test.go b/change_test.go
index 4abea60..383bff6 100644
--- a/change_test.go
+++ b/change_test.go
@@ -80,6 +80,10 @@
 			config.Set(configPollInterval, 10*time.Millisecond)
 		}
 
+		AfterEach(func() {
+			restoreContext()
+		})
+
 		It("test change agent with authorization failure", func() {
 			log.Debug("test change agent with authorization failure")
 			testTokenManager := &dummyTokenManager{make(chan bool)}
@@ -94,7 +98,6 @@
 			<-testTokenManager.invalidateChan
 			log.Debug("closing")
 			<-apidChangeManager.close()
-			restoreContext()
 		})
 
 		It("test change agent with too old snapshot", func() {
@@ -113,7 +116,6 @@
 			<-testSnapshotManager.downloadCalledChan
 			log.Debug("closing")
 			<-apidChangeManager.close()
-			restoreContext()
 		})
 
 		It("change agent should retry with authorization failure", func(done Done) {
@@ -135,7 +137,6 @@
 					go func() {
 						// when close done, all handlers for the first snapshot have been executed
 						<-closeDone
-						restoreContext()
 						close(done)
 					}()
 
@@ -145,7 +146,7 @@
 			apidChangeManager.pollChangeWithBackoff()
 			// auth check fails
 			<-testTokenManager.invalidateChan
-		})
+		}, 2)
 
 	})
 })