Tests and related tweaks

robot.go

@@ -7,6 +7,8 @@ import (
 	"os/signal"
 	"sync/atomic"
 
+	"sync"
+
 	multierror "github.com/hashicorp/go-multierror"
 )
 
@@ -43,15 +45,17 @@ func NewJSONRobot(robot *Robot) *JSONRobot {
 // It contains its own work routine and a collection of
 // custom commands to control a robot remotely via the Gobot api.
 type Robot struct {
-	Name         string
-	Work         func()
-	connections  *Connections
-	devices      *Devices
-	trap         func(chan os.Signal)
-	AutoRun      bool
-	running      atomic.Value
-	done         chan bool
-	workRegistry *RobotWorkRegistry
+	Name               string
+	Work               func()
+	connections        *Connections
+	devices            *Devices
+	trap               func(chan os.Signal)
+	AutoRun            bool
+	running            atomic.Value
+	done               chan bool
+	workRegistry       *RobotWorkRegistry
+	WorkEveryWaitGroup *sync.WaitGroup
+	WorkAfterWaitGroup *sync.WaitGroup
 	Commander
 	Eventer
 }
@@ -143,6 +147,8 @@ func NewRobot(v ...interface{}) *Robot {
 	r.workRegistry = &RobotWorkRegistry{
 		r: make(map[string]*RobotWork),
 	}
+	r.WorkAfterWaitGroup = &sync.WaitGroup{}
+	r.WorkEveryWaitGroup = &sync.WaitGroup{}
 
 	r.running.Store(false)
 	log.Println("Robot", r.Name, "initialized.")

robot_work.go

@@ -1,6 +1,3 @@
-// RobotWork and the RobotWork registry represent units of executing computation
-// managed at the Robot level. Unlike the utility functions gobot.After and gobot.Every,
-// RobotWork units require a context.Context, and can be cancelled externally by calling code.
 package gobot
 
 import (
@@ -25,9 +22,28 @@ const (
 	AfterWorkKind = "after"
 )
 
-// RobotWork represents a unit of work (in the form of an arbitrary Go function)
-// to be done once or on a recurring basis. It encapsulations notions of duration,
-// context, count of successful runs, etc.
+// RobotWork and the RobotWork registry represent units of executing computation
+// managed at the Robot level. Unlike the utility functions gobot.After and gobot.Every,
+// RobotWork units require a context.Context, and can be cancelled externally by calling code.
+//
+// Usage:
+//
+//	someWork := myRobot.Every(context.Background(), time.Second * 2, func(){
+//		fmt.Println("Here I am doing work")
+// 	})
+//
+//	someWork.CallCancelFunc() // Cancel next tick and remove from work registry
+//
+// goroutines for Every and After are run on their own WaitGroups for synchronization:
+//
+//	someWork2 := myRobot.Every(context.Background(), time.Second * 2, func(){
+//		fmt.Println("Here I am doing more work")
+// 	})
+//
+//	somework2.CallCancelFunc()
+//
+//	// wait for both Every calls to finish
+//	robot.WorkEveryWaitGroup().Wait()
 type RobotWork struct {
 	id         uuid.UUID
 	kind       string
@@ -49,6 +65,11 @@ func (rw *RobotWork) CancelFunc() context.CancelFunc {
 	return rw.cancelFunc
 }
 
+// CallCancelFunc calls the context.CancelFunc used to cancel the work
+func (rw *RobotWork) CallCancelFunc() {
+	rw.cancelFunc()
+}
+
 // Ticker returns the time.Ticker used in an Every so that calling code can sync on the same channel
 func (rw *RobotWork) Ticker() *time.Ticker {
 	if rw.kind == AfterWorkKind {
@@ -57,6 +78,11 @@ func (rw *RobotWork) Ticker() *time.Ticker {
 	return rw.ticker
 }
 
+// TickCount returns the number of times the function successfully ran
+func (rw *RobotWork) TickCount() int {
+	return rw.tickCount
+}
+
 // Duration returns the timeout until an After fires or the period of an Every
 func (rw *RobotWork) Duration() time.Duration {
 	return rw.duration
@@ -79,6 +105,7 @@ func (r *Robot) WorkRegistry() *RobotWorkRegistry {
 // Every calls the given function for every tick of the provided duration.
 func (r *Robot) Every(ctx context.Context, d time.Duration, f func()) *RobotWork {
 	rw := r.workRegistry.registerEvery(ctx, d, f)
+	r.WorkEveryWaitGroup.Add(1)
 	go func() {
 	EVERYWORK:
 		for {
@@ -88,19 +115,20 @@ func (r *Robot) Every(ctx context.Context, d time.Duration, f func()) *RobotWork
 				rw.ticker.Stop()
 				break EVERYWORK
 			case <-rw.ticker.C:
-				rw.tickCount++
 				f()
+				rw.tickCount++
 			}
 		}
+		r.WorkEveryWaitGroup.Done()
 	}()
-
 	return rw
 }
 
 // After calls the given function after the provided duration has elapsed
 func (r *Robot) After(ctx context.Context, d time.Duration, f func()) *RobotWork {
 	rw := r.workRegistry.registerAfter(ctx, d, f)
 	ch := time.After(d)
+	r.WorkAfterWaitGroup.Add(1)
 	go func() {
 	AFTERWORK:
 		for {
@@ -112,6 +140,7 @@ func (r *Robot) After(ctx context.Context, d time.Duration, f func()) *RobotWork
 				f()
 			}
 		}
+		r.WorkAfterWaitGroup.Done()
 	}()
 	return rw
 }

robot_work_test.go

@@ -57,6 +57,40 @@ func TestRobotWorkRegistry(t *testing.T) {
 	})
 }
 
+func TestRobotAutomationFunctions(t *testing.T) {
+	t.Run("Every with cancel", func(t *testing.T) {
+		robot := NewRobot("testbot")
+
+		rw := robot.Every(context.Background(), time.Millisecond*10, func() {
+			_ = 1 + 1 // perform mindless computation!
+		})
+
+		time.Sleep(time.Millisecond * 25)
+		rw.CallCancelFunc()
+
+		robot.WorkEveryWaitGroup.Wait()
+
+		assert.Equal(t, 2, rw.tickCount)
+		postDeleteKeys := collectStringKeysFromWorkRegistry(robot.workRegistry)
+		assert.NotContains(t, postDeleteKeys, rw.id.String())
+	})
+
+	t.Run("After with cancel", func(t *testing.T) {
+		robot := NewRobot("testbot")
+
+		rw := robot.After(context.Background(), time.Millisecond*10, func() {
+			_ = 1 + 1 // perform mindless computation!
+		})
+
+		rw.CallCancelFunc()
+
+		robot.WorkAfterWaitGroup.Wait()
+
+		postDeleteKeys := collectStringKeysFromWorkRegistry(robot.workRegistry)
+		assert.NotContains(t, postDeleteKeys, rw.id.String())
+	})
+}
+
 func collectStringKeysFromWorkRegistry(rwr *RobotWorkRegistry) []string {
 	keys := make([]string, len(rwr.r))
 	for k, _ := range rwr.r {