mock: Improvements

Use testing.TB.Cleanup to register Stop inside Start, and also use
slightly better context cancellation check.

echo_test.go

@@ -15,8 +15,7 @@ func TestEcho(t *testing.T) {
 
 	const timeout = time.Millisecond * 200
 
-	service, device := mock.StartService(t)
-	defer service.Stop()
+	_, device := mock.StartService(t)
 
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()

label_test.go

@@ -98,7 +98,6 @@ func TestGetLabel(t *testing.T) {
 	expected.Set("foo")
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 	service.RawStateLabelPayload = &lifxlan.RawStateLabelPayload{
 		Label: expected,
 	}

light/color_test.go

@@ -53,7 +53,6 @@ func TestSetColor(t *testing.T) {
 	label.Set("foo")
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 	service.RawStatePayload = &light.RawStatePayload{
 		Label: label,
 	}

light/wrap_test.go

@@ -21,7 +21,6 @@ func TestWrap(t *testing.T) {
 	label.Set("foo")
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 
 	t.Run(
 		"Normal",

mock/doc_test.go

@@ -29,7 +29,6 @@ func Example_testGetLabel() {
 			expected.Set("foo")
 
 			service, device := mock.StartService(t)
-			defer service.Stop()
 			// This is the payload to be returned by the mock service.
 			service.RawStateLabelPayload = &lifxlan.RawStateLabelPayload{
 				Label: expected,
@@ -64,7 +63,6 @@ func Example_testGetLabelWithHandlerFunc() {
 			expected.Set("foo")
 
 			service, device := mock.StartService(t)
-			defer service.Stop()
 
 			// This defines the handler for GetLabel messages.
 			service.Handlers[lifxlan.GetLabel] = func(
@@ -112,7 +110,6 @@ func Example_testNotEnoughAcks() {
 			const timeout = time.Millisecond * 200
 
 			service, device := mock.StartService(t)
-			defer service.Stop()
 
 			rawTile1 := tile.RawTileDevice{
 				Width:  8,

mock/mock.go

@@ -210,6 +210,8 @@ func StartService(tb testing.TB) (*Service, lifxlan.Device) {
 }
 
 // Start starts the service and returns the device.
+//
+// It also register Stop to TB's Cleanup.
 func (s *Service) Start() lifxlan.Device {
 	s.TB.Helper()
 
@@ -225,6 +227,7 @@ func (s *Service) Start() lifxlan.Device {
 	s.wg.Add(1)
 	go s.handler(conn)
 
+	s.TB.Cleanup(s.Stop)
 	return lifxlan.NewDevice(
 		conn.LocalAddr().String(),
 		lifxlan.ServiceUDP,
@@ -250,9 +253,9 @@ func (s *Service) Reply(
 	message lifxlan.MessageType,
 	payload []byte,
 ) {
-	select {
-	default:
-	case <-s.Context.Done():
+	s.TB.Helper()
+
+	if s.Context.Err() != nil {
 		return
 	}
 
@@ -266,24 +269,22 @@ func (s *Service) Reply(
 		payload,
 	)
 	if err != nil {
-		s.TB.Log(err)
+		s.TB.Logf("lifxlan/mock.Reply: Failed to generate replay message: %v", err)
 		return
 	}
 
-	select {
-	default:
-	case <-s.Context.Done():
+	if s.Context.Err() != nil {
 		return
 	}
 
 	n, err := conn.WriteTo(msg, addr)
 	if err != nil {
-		s.TB.Log(err)
+		s.TB.Logf("lifxlan/mock.Reply: Failed to write replay message: %v", err)
 		return
 	}
 	if n < len(msg) {
 		s.TB.Logf(
-			"lifxlan/mock.Reply: only wrote %d out of %d bytes",
+			"lifxlan/mock.Reply: Only wrote %d out of %d bytes",
 			n,
 			len(msg),
 		)
@@ -296,9 +297,7 @@ func (s *Service) handler(conn net.PacketConn) {
 
 	buf := make([]byte, lifxlan.ResponseReadBufferSize)
 	for {
-		select {
-		default:
-		case <-s.Context.Done():
+		if s.Context.Err() != nil {
 			s.TB.Log(s.Context.Err())
 			return
 		}

power_test.go

@@ -88,7 +88,6 @@ func TestGetPower(t *testing.T) {
 	expected := lifxlan.PowerOn
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 	service.RawStatePowerPayload = &lifxlan.RawStatePowerPayload{
 		Level: expected,
 	}
@@ -121,7 +120,6 @@ func TestSetPower(t *testing.T) {
 	expected := lifxlan.PowerOff
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 
 	var called bool
 

tile/color_test.go

@@ -234,7 +234,6 @@ func TestColorsAPIs(t *testing.T) {
 	label.Set("foo")
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 	service.RawStatePayload = &light.RawStatePayload{
 		Label: label,
 	}

tile/wrap_test.go

@@ -52,7 +52,6 @@ func TestWrap(t *testing.T) {
 	label.Set("foo")
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 	service.RawStatePayload = &light.RawStatePayload{
 		Label: label,
 	}

version_test.go

@@ -101,7 +101,6 @@ func TestGetHardwareVersion(t *testing.T) {
 	}
 
 	service, device := mock.StartService(t)
-	defer service.Stop()
 	service.RawStateVersionPayload = &lifxlan.RawStateVersionPayload{
 		Version: expected,
 	}