Add unit tests for workunitbase and wrap fsnotify library calls

Author: AaronH88

pkg/workceptor/command.go

@@ -219,7 +219,7 @@ func (cw *commandUnit) runCommand(cmd *exec.Cmd) error {
 		})
 	}()
 	go cmdWaiter(cmd, doneChan)
-	go cw.monitorLocalStatus()
+	go cw.MonitorLocalStatus()
 
 	return nil
 }
@@ -263,7 +263,7 @@ func (cw *commandUnit) Restart() error {
 		// Job never started - mark it failed
 		cw.UpdateBasicStatus(WorkStateFailed, "Pending at restart", stdoutSize(cw.UnitDir()))
 	}
-	go cw.monitorLocalStatus()
+	go cw.MonitorLocalStatus()
 
 	return nil
 }
@@ -331,7 +331,7 @@ func (cfg CommandWorkerCfg) NewWorker(w *Workceptor, unitID string, workType str
 		baseParams:         cfg.Params,
 		allowRuntimeParams: cfg.AllowRuntimeParams,
 	}
-	cw.BaseWorkUnit.Init(w, unitID, workType)
+	cw.BaseWorkUnit.Init(w, unitID, workType, FileSystem{}, nil)
 
 	return cw
 }

pkg/workceptor/json_test.go

@@ -22,7 +22,7 @@ func newCommandWorker(w *Workceptor, unitID string, workType string) WorkUnit {
 		baseParams:         "foo",
 		allowRuntimeParams: true,
 	}
-	cw.BaseWorkUnit.Init(w, unitID, workType)
+	cw.BaseWorkUnit.Init(w, unitID, workType, FileSystem{}, nil)
 
 	return cw
 }

pkg/workceptor/kubernetes.go

@@ -1272,7 +1272,7 @@ func (kw *kubeUnit) startOrRestart() error {
 	} else {
 		go kw.runWorkUsingLogger()
 	}
-	go kw.monitorLocalStatus()
+	go kw.MonitorLocalStatus()
 
 	return nil
 }
@@ -1388,7 +1388,7 @@ func (cfg KubeWorkerCfg) NewWorker(w *Workceptor, unitID string, workType string
 		deletePodOnRestart:  cfg.DeletePodOnRestart,
 		namePrefix:          fmt.Sprintf("%s-", strings.ToLower(cfg.WorkType)),
 	}
-	ku.BaseWorkUnit.Init(w, unitID, workType)
+	ku.BaseWorkUnit.Init(w, unitID, workType, FileSystem{}, nil)
 
 	return ku
 }

pkg/workceptor/mock_workceptor/stdio_utils.go

@@ -63,7 +63,7 @@ func (cfg workPythonCfg) NewWorker(w *Workceptor, unitID string, workType string
 		function: cfg.Function,
 		config:   cfg.Config,
 	}
-	cw.BaseWorkUnit.Init(w, unitID, workType)
+	cw.BaseWorkUnit.Init(w, unitID, workType, FileSystem{}, nil)
 
 	return cw
 }

pkg/workceptor/mock_workceptor/workunitbase.go

@@ -680,7 +680,7 @@ func (rw *remoteUnit) Release(force bool) error {
 
 func newRemoteWorker(w *Workceptor, unitID, workType string) WorkUnit {
 	rw := &remoteUnit{logger: w.nc.GetLogger()}
-	rw.BaseWorkUnit.Init(w, unitID, workType)
+	rw.BaseWorkUnit.Init(w, unitID, workType, FileSystem{}, nil)
 	red := &remoteExtraData{}
 	red.RemoteParams = make(map[string]string)
 	rw.status.ExtraData = red

pkg/workceptor/python.go

@@ -16,6 +16,7 @@ type FileSystemer interface {
 	OpenFile(name string, flag int, perm os.FileMode) (*os.File, error)
 	Stat(name string) (os.FileInfo, error)
 	Open(name string) (*os.File, error)
+	RemoveAll(path string) error
 }
 
 // FileSystem represents the real filesystem.
@@ -36,6 +37,11 @@ func (FileSystem) Open(name string) (*os.File, error) {
 	return os.Open(name)
 }
 
+// RemoveAll removes path and any children it contains.
+func (FileSystem) RemoveAll(path string) error {
+	return os.RemoveAll(path)
+}
+
 // FileWriteCloser wraps io.WriteCloser.
 type FileWriteCloser interface {
 	io.WriteCloser

pkg/workceptor/remote_work.go

@@ -32,6 +32,28 @@ const (
 	WorkStateCanceled  = 4
 )
 
+type WatcherWrapper interface {
+	Add(name string) error
+	Close() error
+	EventChannel() chan fsnotify.Event
+}
+
+type RealWatcher struct {
+	watcher *fsnotify.Watcher
+}
+
+func (rw *RealWatcher) Add(name string) error {
+	return rw.watcher.Add(name)
+}
+
+func (rw *RealWatcher) Close() error {
+	return rw.watcher.Close()
+}
+
+func (rw *RealWatcher) EventChannel() chan fsnotify.Event {
+	return rw.watcher.Events
+}
+
 // IsComplete returns true if a given WorkState indicates the job is finished.
 func IsComplete(workState int) bool {
 	return workState == WorkStateSucceeded || workState == WorkStateFailed
@@ -74,10 +96,12 @@ type BaseWorkUnit struct {
 	lastUpdateErrorLock *sync.RWMutex
 	ctx                 context.Context
 	cancel              context.CancelFunc
+	fs                  FileSystemer
+	watcher             WatcherWrapper
 }
 
 // Init initializes the basic work unit data, in memory only.
-func (bwu *BaseWorkUnit) Init(w *Workceptor, unitID string, workType string) {
+func (bwu *BaseWorkUnit) Init(w *Workceptor, unitID string, workType string, fs FileSystemer, watcher WatcherWrapper) {
 	bwu.w = w
 	bwu.status.State = WorkStatePending
 	bwu.status.Detail = "Unit Created"
@@ -90,6 +114,17 @@ func (bwu *BaseWorkUnit) Init(w *Workceptor, unitID string, workType string) {
 	bwu.statusLock = &sync.RWMutex{}
 	bwu.lastUpdateErrorLock = &sync.RWMutex{}
 	bwu.ctx, bwu.cancel = context.WithCancel(w.ctx)
+	bwu.fs = fs
+	if watcher != nil {
+		bwu.watcher = watcher
+	} else {
+		watcher, err := fsnotify.NewWatcher()
+		if err == nil {
+			bwu.watcher = &RealWatcher{watcher: watcher}
+		} else {
+			bwu.watcher = nil
+		}
+	}
 }
 
 // Error logs message with unitID prepended.
@@ -340,32 +375,31 @@ func (bwu *BaseWorkUnit) LastUpdateError() error {
 	return bwu.lastUpdateError
 }
 
-// monitorLocalStatus watches a unit dir and keeps the in-memory workUnit up to date with status changes.
-func (bwu *BaseWorkUnit) monitorLocalStatus() {
+// MonitorLocalStatus watches a unit dir and keeps the in-memory workUnit up to date with status changes.
+func (bwu *BaseWorkUnit) MonitorLocalStatus() {
 	statusFile := path.Join(bwu.UnitDir(), "status")
-	watcher, err := fsnotify.NewWatcher()
-	if err == nil {
-		err = watcher.Add(statusFile)
+	if bwu.watcher != nil {
+		err := bwu.watcher.Add(statusFile)
 		if err == nil {
 			defer func() {
-				_ = watcher.Close()
+				_ = bwu.watcher.Close()
 			}()
 		} else {
-			_ = watcher.Close()
-			watcher = nil
+			_ = bwu.watcher.Close()
+			bwu.watcher = nil
 		}
 	} else {
-		watcher = nil
+		bwu.watcher = nil
 	}
-	fi, err := os.Stat(statusFile)
+	fi, err := bwu.fs.Stat(statusFile)
 	if err != nil {
 		fi = nil
 	}
 	var watcherEvents chan fsnotify.Event
-	if watcher == nil {
+	if bwu.watcher == nil {
 		watcherEvents = make(chan fsnotify.Event)
 	} else {
-		watcherEvents = watcher.Events
+		watcherEvents = bwu.watcher.EventChannel()
 	}
 loop:
 	for {
@@ -380,7 +414,7 @@ loop:
 				}
 			}
 		case <-time.After(time.Second):
-			newFi, err := os.Stat(statusFile)
+			newFi, err := bwu.fs.Stat(statusFile)
 			if err == nil {
 				if fi == nil || fi.ModTime() != newFi.ModTime() {
 					fi = newFi
@@ -425,7 +459,7 @@ func (bwu *BaseWorkUnit) Release(force bool) error {
 	defer bwu.statusLock.Unlock()
 	attemptsLeft := 3
 	for {
-		err := os.RemoveAll(bwu.UnitDir())
+		err := bwu.fs.RemoveAll(bwu.UnitDir())
 		if force {
 			break
 		} else if err != nil {
@@ -451,11 +485,15 @@ func (bwu *BaseWorkUnit) Release(force bool) error {
 	return nil
 }
 
+func (bwu *BaseWorkUnit) CancelContext() {
+	bwu.cancel()
+}
+
 // =============================================================================================== //
 
 func newUnknownWorker(w *Workceptor, unitID string, workType string) WorkUnit {
 	uu := &unknownUnit{}
-	uu.BaseWorkUnit.Init(w, unitID, workType)
+	uu.BaseWorkUnit.Init(w, unitID, workType, FileSystem{}, nil)
 
 	return uu
 }