Merge pull request #477 from spiral/feature/worker_watcher_testsv2.0.0-beta8

feat(worker_watcher): Add stack tests

1 files changed, 10 insertions, 8 deletions

diff --git a/pkg/worker_watcher/stack.go b/pkg/worker_watcher/stack.go
index 55f1f52a..788750dc 100644
--- a/pkg/worker_watcher/stack.go
+++ b/pkg/worker_watcher/stack.go
@@ -10,17 +10,19 @@ import (
 )
 
 type Stack struct {
-	workers            []worker.BaseProcess
-	mutex              sync.RWMutex
-	destroy            bool
-	actualNumOfWorkers int64
+	workers             []worker.BaseProcess
+	mutex               sync.RWMutex
+	destroy             bool
+	actualNumOfWorkers  uint64
+	initialNumOfWorkers uint64
 }
 
-func NewWorkersStack() *Stack {
+func NewWorkersStack(initialNumOfWorkers uint64) *Stack {
 	w := runtime.NumCPU()
 	return &Stack{
-		workers:            make([]worker.BaseProcess, 0, w),
-		actualNumOfWorkers: 0,
+		workers:             make([]worker.BaseProcess, 0, w),
+		actualNumOfWorkers:  0,
+		initialNumOfWorkers: initialNumOfWorkers,
 	}
 }
 
@@ -113,7 +115,7 @@ func (stack *Stack) Destroy(ctx context.Context) {
 		case <-tt.C:
 			stack.mutex.Lock()
 			// that might be one of the workers is working
-			if len(stack.workers) != int(stack.actualNumOfWorkers) {
+			if stack.initialNumOfWorkers != stack.actualNumOfWorkers {
 				stack.mutex.Unlock()
 				continue
 			}