From e7c6b8e7a50de133f0373c5217cdc0768bc40d10 Mon Sep 17 00:00:00 2001
From: Yousong Zhou <zhouyousong@yunionyun.com>
Date: Wed, 19 Dec 2018 10:13:23 +0000
Subject: [PATCH] appsrv: fix TestWorkerManagerError()

---
 pkg/appsrv/workers.go      | 16 ------------
 pkg/appsrv/workers_test.go | 51 ++++++++++++++++++++++++++------------
 2 files changed, 35 insertions(+), 32 deletions(-)

diff --git a/pkg/appsrv/workers.go b/pkg/appsrv/workers.go
index a751067466..3a8d7076f9 100644
--- a/pkg/appsrv/workers.go
+++ b/pkg/appsrv/workers.go
@@ -247,19 +247,3 @@ func WorkerStatsHandler(ctx context.Context, w http.ResponseWriter, r *http.Requ
 	result.Add(jsonutils.Marshal(&stats), "workers")
 	fmt.Fprintf(w, result.String())
 }
-
-func WaitChannel(ch chan interface{}) interface{} {
-	var ret interface{}
-	stop := false
-	for !stop {
-		select {
-		case c, more := <-ch:
-			if more {
-				ret = c
-			} else {
-				stop = true
-			}
-		}
-	}
-	return ret
-}
diff --git a/pkg/appsrv/workers_test.go b/pkg/appsrv/workers_test.go
index bc1527ab37..d02cf762c1 100644
--- a/pkg/appsrv/workers_test.go
+++ b/pkg/appsrv/workers_test.go
@@ -1,6 +1,7 @@
 package appsrv
 
 import (
+	"sync"
 	"testing"
 	"time"
 )
@@ -27,21 +28,39 @@ func TestWorkerManager(t *testing.T) {
 
 func TestWorkerManagerError(t *testing.T) {
 	wm := NewWorkerManager("testwm", 2, 10)
-	err := make(chan interface{})
-	wm.Run(func() {
-		panic("Panic inside worker")
-	}, nil, err)
-	e := WaitChannel(err)
-	if e == nil {
-		t.Error("Panic not captured")
+	errCbFactory := func(wg *sync.WaitGroup, errMark *bool) func(error) {
+		return func(error) {
+			defer wg.Done()
+			if errMark != nil && !*errMark {
+				*errMark = true
+			}
+		}
 	}
-	err = make(chan interface{})
-	wm.Run(func() {
-		time.Sleep(1 * time.Second)
-	}, nil, err)
-	e = WaitChannel(err)
-	if e != nil {
-		t.Error("Should no error")
-	}
-
+	t.Run("normal", func(t *testing.T) {
+		wg := &sync.WaitGroup{}
+		errMark := false
+		errCb := errCbFactory(wg, &errMark)
+		wg.Add(1)
+		wm.Run(func() {
+			defer wg.Done()
+		}, nil, errCb)
+		wg.Wait()
+		if errMark {
+			t.Errorf("should be normal")
+		}
+	})
+	t.Run("panic", func(t *testing.T) {
+		wg := &sync.WaitGroup{}
+		errMark := false
+		errCb := errCbFactory(wg, &errMark)
+		wg.Add(2) // 1 for errCb
+		wm.Run(func() {
+			defer wg.Done()
+			panic("panic inside worker")
+		}, nil, errCb)
+		wg.Wait()
+		if !errMark {
+			t.Errorf("expecting error")
+		}
+	})
 }