[K10-2154] Add cancelFunc to testutil (#5358)

* Add cancelFunc to testutil

* Refactor tests

* Remove unused var

pkg/testutil/func.go

@@ -14,6 +14,7 @@ const (
 	WaitFuncName   = "WaitFunc"
 	ArgFuncName    = "ArgFunc"
 	OutputFuncName = "OutputFunc"
+	CancelFuncName = "CancelFunc"
 )
 
 var (
@@ -33,6 +34,7 @@ func waitFunc(context.Context, param.TemplateParams, map[string]interface{}) (ma
 	<-waitFuncCh
 	return nil, nil
 }
+
 func argsFunc(ctx context.Context, tp param.TemplateParams, args map[string]interface{}) (map[string]interface{}, error) {
 	argFuncCh <- args
 	return nil, nil
@@ -43,6 +45,11 @@ func outputFunc(ctx context.Context, tp param.TemplateParams, args map[string]in
 	return args, nil
 }
 
+func cancelFunc(ctx context.Context, tp param.TemplateParams, args map[string]interface{}) (map[string]interface{}, error) {
+	<-ctx.Done()
+	return nil, ctx.Err()
+}
+
 func init() {
 	failFuncCh = make(chan error)
 	waitFuncCh = make(chan struct{})
@@ -52,6 +59,7 @@ func init() {
 	registerMockKanisterFunc(WaitFuncName, waitFunc)
 	registerMockKanisterFunc(ArgFuncName, argsFunc)
 	registerMockKanisterFunc(OutputFuncName, outputFunc)
+	registerMockKanisterFunc(CancelFuncName, cancelFunc)
 }
 
 func registerMockKanisterFunc(name string, f func(context.Context, param.TemplateParams, map[string]interface{}) (map[string]interface{}, error)) {

pkg/testutil/func_test.go

@@ -2,6 +2,7 @@ package testutil
 
 import (
 	"context"
+	"strings"
 
 	. "gopkg.in/check.v1"
 
@@ -65,3 +66,21 @@ func (s *FuncSuite) TestOutputFunc(c *C) {
 	}()
 	c.Assert(OutputFuncOut(), DeepEquals, args)
 }
+
+func (s *FuncSuite) TestCancelFunc(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan bool)
+	go func() {
+		_, err := cancelFunc(ctx, param.TemplateParams{}, nil)
+		c.Assert(err, NotNil)
+		c.Assert(strings.Contains(err.Error(), "context canceled"), Equals, true)
+		close(done)
+	}()
+	select {
+	case <-done:
+		c.FailNow()
+	default:
+	}
+	cancel()
+	<-done
+}