From 700a5b0cdc23ae800567d5cec0f8c9b1d11980b7 Mon Sep 17 00:00:00 2001 From: Yuxuan 'fishy' Wang Date: Sat, 23 Apr 2022 21:23:57 -0700 Subject: [PATCH] mock: Improvements Use testing.TB.Cleanup to register Stop inside Start, and also use slightly better context cancellation check. --- echo_test.go | 3 +-- label_test.go | 1 - light/color_test.go | 1 - light/wrap_test.go | 1 - mock/doc_test.go | 3 --- mock/mock.go | 23 +++++++++++------------ power_test.go | 2 -- tile/color_test.go | 1 - tile/wrap_test.go | 1 - version_test.go | 1 - 10 files changed, 12 insertions(+), 25 deletions(-) diff --git a/echo_test.go b/echo_test.go index 900b2f2..98acc01 100644 --- a/echo_test.go +++ b/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() diff --git a/label_test.go b/label_test.go index 03b1e99..106661e 100644 --- a/label_test.go +++ b/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, } diff --git a/light/color_test.go b/light/color_test.go index cb81eec..0676829 100644 --- a/light/color_test.go +++ b/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, } diff --git a/light/wrap_test.go b/light/wrap_test.go index e4c8fda..fc0942e 100644 --- a/light/wrap_test.go +++ b/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", diff --git a/mock/doc_test.go b/mock/doc_test.go index db3c3be..6d38ac6 100644 --- a/mock/doc_test.go +++ b/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, diff --git a/mock/mock.go b/mock/mock.go index cff3d0a..b603f82 100644 --- a/mock/mock.go +++ b/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 } diff --git a/power_test.go b/power_test.go index 0198c8a..da16a1a 100644 --- a/power_test.go +++ b/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 diff --git a/tile/color_test.go b/tile/color_test.go index abbf07f..ed0cfa2 100644 --- a/tile/color_test.go +++ b/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, } diff --git a/tile/wrap_test.go b/tile/wrap_test.go index 6d2fc16..accd85b 100644 --- a/tile/wrap_test.go +++ b/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, } diff --git a/version_test.go b/version_test.go index daa3a2f..00d6e3c 100644 --- a/version_test.go +++ b/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, }