as7262: add internal context.context (#399)

Add internal context.context
Fixes some fragile time based tests.
Halt() is now synchronous, will only return if the context has been cancelled.
Remove go routine leak in New().
Fix spelling.

experimental/devices/as7262/as7262.go

@@ -5,6 +5,7 @@
 package as7262
 
 import (
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -33,29 +34,35 @@ var DefaultOpts = Opts{
 // New opens a handle to an AS7262 sensor.
 func New(bus i2c.Bus, opts *Opts) (*Dev, error) {
 	// The nil or zero values for gain and interrupt are valid
+	c := make(chan struct{})
+	close(c)
 	return &Dev{
 		c:         &i2c.Dev{Bus: bus, Addr: 0x49},
 		gain:      opts.Gain,
 		interrupt: opts.InterruptPin,
-		done:      make(chan struct{}, 1),
-		timeout:   200 * time.Millisecond,
+		cancel:    func() {},
+		done:      c,
 	}, nil
 }
 
+var sensorTimeout = 200 * time.Millisecond
+
+// waitForSensor is overridden in tests.
+var waitForSensor = time.After
+
 // Dev is a handle to the as7262 sensor.
 type Dev struct {
 	c         conn.Conn
-	timeout   time.Duration
 	interrupt gpio.PinIn
 
 	// Mutable
 	mu   sync.Mutex
 	gain Gain
-	done chan struct{}
 
-	// Guards canceled
-	canceledMu sync.Mutex
-	canceled   bool
+	// cancelMu guards cancel and done.
+	cancelMu sync.Mutex
+	cancel   context.CancelFunc
+	done     chan struct{}
 }
 
 // Spectrum is the reading from the sensor including the actual sensor state for
@@ -123,25 +130,30 @@ func (b Band) String() string {
 func (d *Dev) Sense(ledDrive physic.ElectricCurrent, senseTime time.Duration) (Spectrum, error) {
 
 	d.mu.Lock()
-	defer d.mu.Unlock()
-
-	done := make(chan struct{}, 1)
-	d.canceledMu.Lock()
-	d.done = done
-	d.canceled = false
-	d.canceledMu.Unlock()
 
 	it, integration := calcSenseTime(senseTime)
-	if err := d.writeVirtualRegister(intergrationReg, it); err != nil {
+	ctx, cancel := context.WithCancel(context.Background())
+	d.cancelMu.Lock()
+	d.cancel = cancel
+	d.done = make(chan struct{})
+	d.cancelMu.Unlock()
+
+	defer func() {
+		close(d.done)
+		d.cancel()
+		d.mu.Unlock()
+	}()
+
+	if err := d.writeVirtualRegister(ctx, integrationReg, it); err != nil {
 		return Spectrum{}, err
 	}
 	led, drive := calcLed(ledDrive)
 
-	if err := d.writeVirtualRegister(ledControlReg, led); err != nil {
+	if err := d.writeVirtualRegister(ctx, ledControlReg, led); err != nil {
 		return Spectrum{}, err
 	}
 
-	if err := d.writeVirtualRegister(controlReg, uint8(allOneShot)|uint8(d.gain)); err != nil {
+	if err := d.writeVirtualRegister(ctx, controlReg, uint8(allOneShot)|uint8(d.gain)); err != nil {
 		return Spectrum{}, err
 	}
 
@@ -149,40 +161,40 @@ func (d *Dev) Sense(ledDrive physic.ElectricCurrent, senseTime time.Duration) (S
 		isEdge := make(chan bool)
 		go func() {
 			// TODO(NeuralSpaz): Test on hardware.
-			isEdge <- d.interrupt.WaitForEdge(integration*2 + d.timeout)
+			isEdge <- d.interrupt.WaitForEdge(integration*2 + sensorTimeout)
 		}()
 		select {
 		case edge := <-isEdge:
 			if !edge {
 				return Spectrum{}, errPinTimeout
 			}
-		case <-d.done:
+		case <-ctx.Done():
 			return Spectrum{}, errHalted
 		}
 	} else {
 		select {
 		// WaitForSensor is time.After().
 		case <-waitForSensor(integration * 2):
-			if err := d.pollDataReady(); err != nil {
+			if err := d.pollDataReady(ctx); err != nil {
 				return Spectrum{}, err
 			}
-		case <-d.done:
+		case <-ctx.Done():
 			return Spectrum{}, errHalted
 		}
 
 	}
 
-	if err := d.writeVirtualRegister(ledControlReg, 0x00); err != nil {
+	if err := d.writeVirtualRegister(ctx, ledControlReg, 0x00); err != nil {
 		return Spectrum{}, err
 	}
 
 	raw := make([]byte, 12)
-	if err := d.readVirtualRegister(rawBase, raw); err != nil {
+	if err := d.readVirtualRegister(ctx, rawBase, raw); err != nil {
 		return Spectrum{}, err
 	}
 
 	cal := make([]byte, 24)
-	if err := d.readVirtualRegister(calBase, cal); err != nil {
+	if err := d.readVirtualRegister(ctx, calBase, cal); err != nil {
 		return Spectrum{}, err
 	}
 
@@ -201,7 +213,7 @@ func (d *Dev) Sense(ledDrive physic.ElectricCurrent, senseTime time.Duration) (S
 	rcal := float64(math.Float32frombits(binary.BigEndian.Uint32(cal[20:24])))
 
 	traw := make([]byte, 1)
-	if err := d.readVirtualRegister(deviceTemperatureReg, traw); err != nil {
+	if err := d.readVirtualRegister(ctx, deviceTemperatureReg, traw); err != nil {
 		return Spectrum{}, err
 	}
 	temperature := physic.Temperature(int8(traw[0]))*physic.Kelvin + physic.ZeroCelsius
@@ -221,17 +233,14 @@ func (d *Dev) Sense(ledDrive physic.ElectricCurrent, senseTime time.Duration) (S
 	}, nil
 }
 
-var waitForSensor = time.After
-
 // Halt stops any pending operations. Repeated calls to Halt do nothing.
 func (d *Dev) Halt() error {
-	d.canceledMu.Lock()
-	defer d.canceledMu.Unlock()
-
-	if !d.canceled {
-		d.done <- struct{}{}
-		d.canceled = true
-	}
+	d.cancelMu.Lock()
+	defer d.cancelMu.Unlock()
+	d.cancel()
+	// A receive can always proceed on a closed channel we can use that
+	// to signal that the running process has been canceled correctly.
+	_, _ = <-d.done
 	return nil
 }
 
@@ -284,13 +293,7 @@ func (d *Dev) Gain(gain Gain) error {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
-	done := make(chan struct{}, 1)
-	d.canceledMu.Lock()
-	d.done = done
-	d.canceled = false
-	d.canceledMu.Unlock()
-
-	if err := d.writeVirtualRegister(controlReg, uint8(gain)); err != nil {
+	if err := d.writeVirtualRegister(context.Background(), controlReg, uint8(gain)); err != nil {
 		return err
 	}
 	d.gain = gain
@@ -301,10 +304,10 @@ func (d *Dev) Gain(gain Gain) error {
 // MSB of the register must be set when writing the register to the write
 // register, also status register must be checked for pending writes or data may
 // be discarded.
-func (d *Dev) writeVirtualRegister(register, data byte) error {
+func (d *Dev) writeVirtualRegister(ctx context.Context, register, data byte) error {
 
 	// Check for pending writes.
-	if err := d.pollStatus(writing); err != nil {
+	if err := d.pollStatus(ctx, writing); err != nil {
 		return err
 	}
 
@@ -314,7 +317,7 @@ func (d *Dev) writeVirtualRegister(register, data byte) error {
 	}
 
 	// Check for pending writes again.
-	if err := d.pollStatus(writing); err != nil {
+	if err := d.pollStatus(ctx, writing); err != nil {
 		return err
 	}
 
@@ -324,18 +327,17 @@ func (d *Dev) writeVirtualRegister(register, data byte) error {
 	}
 
 	return nil
-
 }
 
 // AS7262 protocol uses virtual registers. To read a virtual register the
 // pointer to the virtual register must be written to the write register. Status
 // register must be checked for any pending reads or data may be invalid, then
 // data maybe read from the read register.
-func (d *Dev) readVirtualRegister(register byte, data []byte) error {
+func (d *Dev) readVirtualRegister(ctx context.Context, register byte, data []byte) error {
 	rx := make([]byte, 1)
 	for i := 0; i < len(data); i++ {
 		// Check for pending reads.
-		if err := d.pollStatus(clearBuffer); err != nil {
+		if err := d.pollStatus(ctx, clearBuffer); err != nil {
 			return err
 		}
 
@@ -345,7 +347,7 @@ func (d *Dev) readVirtualRegister(register byte, data []byte) error {
 		}
 
 		// Check if read buffer is ready.
-		if err := d.pollStatus(reading); err != nil {
+		if err := d.pollStatus(ctx, reading); err != nil {
 			return err
 		}
 
@@ -360,12 +362,12 @@ func (d *Dev) readVirtualRegister(register byte, data []byte) error {
 }
 
 // Polls the data ready bit in the control register(virtual)
-func (d *Dev) pollDataReady() error {
-	timeout := time.NewTimer(d.timeout)
+func (d *Dev) pollDataReady(ctx context.Context) error {
+	timeout := time.NewTimer(sensorTimeout)
 	defer timeout.Stop()
 
 	for {
-		if err := d.pollStatus(clearBuffer); err != nil {
+		if err := d.pollStatus(ctx, clearBuffer); err != nil {
 			return err
 		}
 
@@ -375,7 +377,7 @@ func (d *Dev) pollDataReady() error {
 		}
 
 		// Check if read buffer is ready.
-		if err := d.pollStatus(reading); err != nil {
+		if err := d.pollStatus(ctx, reading); err != nil {
 			return err
 		}
 
@@ -394,11 +396,10 @@ func (d *Dev) pollDataReady() error {
 		case <-timeout.C:
 			// Return error if it takes too long.
 			return errStatusDeadline
-		case <-d.done:
+		case <-ctx.Done():
 			return errHalted
 		}
 	}
-
 }
 
 type direction byte
@@ -416,19 +417,18 @@ const (
 // provides a way to repeatedly check if there are any pending reads or writes
 // in the relevent buffer before a transaction while with a timeout.
 // Direction is used to set which buffer is being polled to be ready.
-func (d *Dev) pollStatus(dir direction) error {
-	timeout := time.NewTimer(d.timeout)
+func (d *Dev) pollStatus(ctx context.Context, dir direction) error {
+	timeout := time.NewTimer(sensorTimeout)
 	defer timeout.Stop()
-	// Check if already canceled first
+
 	select {
-	case <-d.done:
+	case <-ctx.Done():
 		return errHalted
 	default:
-		// Proceed.
 	}
+
 	status := make([]byte, 1)
 	for {
-
 		// Read status register.
 		err := d.c.Tx([]byte{statusReg}, status)
 		if err != nil {
@@ -462,13 +462,14 @@ func (d *Dev) pollStatus(dir direction) error {
 				return nil
 			}
 		}
+
 		select {
 		case <-time.After(5 * time.Millisecond):
 			// Polling interval.
 		case <-timeout.C:
 			// Return error if it takes too long.
 			return errStatusDeadline
-		case <-d.done:
+		case <-ctx.Done():
 			return errHalted
 		}
 	}
@@ -551,7 +552,7 @@ const (
 	hardwareVersion      = 0x00
 	firmwareVersion      = 0x02
 	controlReg           = 0x04
-	intergrationReg      = 0x05
+	integrationReg       = 0x05
 	deviceTemperatureReg = 0x06
 	ledControlReg        = 0x07
 	// RawBase used as base for reading uint16 values, data must be sequentially.

experimental/devices/as7262/as7262_test.go

@@ -5,6 +5,7 @@
 package as7262
 
 import (
+	"context"
 	"reflect"
 	"testing"
 	"time"
@@ -16,18 +17,16 @@ import (
 	"periph.io/x/periph/conn/physic"
 )
 
+var defaultSensorTimeOut = time.Millisecond * 200
+
 func TestDev_Sense(t *testing.T) {
 
 	type timefunc func(*Dev) func(time.Duration) <-chan time.Time
 
-	defer func() {
-		waitForSensor = time.After
-	}()
-
 	haltit := func(dev *Dev) func(time.Duration) <-chan time.Time {
 		return func(d time.Duration) <-chan time.Time {
 			t := make(chan time.Time, 1)
-			dev.Halt()
+			go dev.Halt()
 			return t
 		}
 	}
@@ -35,7 +34,7 @@ func TestDev_Sense(t *testing.T) {
 		return func(d time.Duration) <-chan time.Time {
 			t := make(chan time.Time, 1)
 			t <- time.Now()
-			dev.timeout = time.Millisecond * 1
+			sensorTimeout = time.Millisecond * 1
 			return t
 		}
 	}
@@ -89,17 +88,6 @@ func TestDev_Sense(t *testing.T) {
 			want:     Spectrum{},
 			tx:       sensorTestCaseInteruptValidRead,
 		},
-		{
-			name: "haltBeforeforEdge",
-			opts: Opts{
-				InterruptPin: intPin,
-			},
-			waiter:   dontwait,
-			sendEdge: false,
-			wantErr:  errHalted,
-			want:     Spectrum{},
-			tx:       sensorTestCaseInteruptValidRead,
-		},
 		{
 			name: "ioErrorWritingIntergrationReg",
 			opts: Opts{
@@ -169,10 +157,19 @@ func TestDev_Sense(t *testing.T) {
 		},
 	}
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			defer func() {
+				// cleanup
+				waitForSensor = time.After
+				sensorTimeout = defaultSensorTimeOut
+			}()
+
 			bus := &i2ctest.Playback{
 				Ops:       tt.tx,
 				DontPanic: true,
 			}
+
 			d, _ := New(bus, &tt.opts)
 
 			waitForSensor = tt.waiter(d)
@@ -180,12 +177,6 @@ func TestDev_Sense(t *testing.T) {
 			if d.interrupt != nil && tt.sendEdge {
 				intPin.EdgesChan <- gpio.High
 			}
-		if tt.name == "haltBeforeforEdge" {
-			// Time must be less than senseTime.
-			time.AfterFunc(time.Millisecond, func() {
-				d.Halt()
-			})
-		}
 
 			got, err := d.Sense(physic.MilliAmpere*100, time.Millisecond*3)
 
@@ -202,6 +193,7 @@ func TestDev_Sense(t *testing.T) {
 			if !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("Dev.Sense() = %v, want %v", got, tt.want)
 			}
+		})
 	}
 }
 
@@ -272,7 +264,7 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{}},
 			},
 			dir:     reading,
-			timeout: time.Millisecond * 1,
+			timeout: time.Second,
 			wantErr: &IOError{"reading status register", nil},
 		},
 		{
@@ -290,8 +282,8 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
 			},
 			dir:     reading,
-			timeout: time.Millisecond * 1000,
-			halt:    time.Nanosecond,
+			timeout: time.Second,
+			halt:    1,
 			wantErr: errHalted,
 		},
 		{
@@ -300,7 +292,7 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
 			},
 			dir:     reading,
-			timeout: time.Millisecond * 1000,
+			timeout: time.Second,
 			wantErr: nil,
 		},
 		{
@@ -309,7 +301,7 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
 			},
 			dir:     writing,
-			timeout: time.Millisecond * 1000,
+			timeout: time.Second,
 			wantErr: nil,
 		},
 		{
@@ -318,7 +310,7 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
 			},
 			dir:     clearBuffer,
-			timeout: time.Millisecond * 1000,
+			timeout: time.Second,
 			wantErr: nil,
 		},
 		{
@@ -329,7 +321,7 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
 			},
 			dir:     clearBuffer,
-			timeout: time.Millisecond * 1000,
+			timeout: time.Second,
 			wantErr: nil,
 		},
 		{
@@ -339,57 +331,41 @@ func TestDev_pollStatus(t *testing.T) {
 				{Addr: 0x49, W: []byte{readReg}, R: []byte{}},
 			},
 			dir:     clearBuffer,
-			timeout: time.Millisecond * 1000,
+			timeout: time.Second,
 			wantErr: &IOError{"clearing buffer", nil},
 		},
-		{
-			name: "retry1",
-			tx: []i2ctest.IO{
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
-			},
-			dir:     reading,
-			timeout: time.Millisecond * 100,
-			wantErr: nil,
-		},
-		{
-			name: "retrywithHalt",
-			tx: []i2ctest.IO{
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
-			},
-			dir:     reading,
-			timeout: time.Millisecond * 100,
-			halt:    time.Millisecond * 3,
-			wantErr: errHalted,
-		},
 	}
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			bus := &i2ctest.Playback{
 				Ops:       tt.tx,
 				DontPanic: true,
 			}
 
+			ctx, cancel := context.WithCancel(context.Background())
+
+			sensorTimeout = tt.timeout
+			defer func() {
+				// cleanup
+				sensorTimeout = defaultSensorTimeOut
+				cancel()
+			}()
+
 			d := &Dev{
 				c:      &i2c.Dev{Bus: bus, Addr: 0x49},
-			done:    make(chan struct{}, 1),
-			timeout: tt.timeout,
+				cancel: cancel,
+				done:   make(chan struct{}),
 			}
-		defer d.Halt()
-		if tt.halt > time.Nanosecond {
+
+			if tt.halt != 0 {
 				go func() {
-				time.Sleep(tt.halt)
-				d.Halt()
+					if err := d.Halt(); err != nil {
+						t.Errorf("Halt() expected nil but got %v", err)
+					}
 				}()
-		} else if tt.halt != 0 {
-			d.Halt()
-			d.Halt()
 			}
 
-		got := d.pollStatus(tt.dir)
-		// t.Errorf("expected error: %v but got: %v", tt.wantErr, got)
+			got := d.pollStatus(ctx, tt.dir)
 
 			if _, ok := tt.wantErr.(*IOError); ok {
 				if _, ok := got.(*IOError); !ok {
@@ -401,6 +377,7 @@ func TestDev_pollStatus(t *testing.T) {
 			} else if got != tt.wantErr {
 				t.Errorf("expected error: %v but got: %v", tt.wantErr, got)
 			}
+		})
 	}
 }
 
@@ -463,27 +440,36 @@ func TestDev_writeVirtualRegister(t *testing.T) {
 		},
 	}
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			bus := &i2ctest.Playback{
 				Ops:       tt.tx,
 				DontPanic: true,
 			}
 
+			ctx, cancel := context.WithCancel(context.Background())
+
+			sensorTimeout = tt.timeout
+			defer func() {
+				// cleanup
+				sensorTimeout = defaultSensorTimeOut
+				cancel()
+			}()
+
 			d := &Dev{
 				c:      &i2c.Dev{Bus: bus, Addr: 0x49},
-			done:    make(chan struct{}, 1),
-			timeout: tt.timeout,
+				cancel: cancel,
+				done:   make(chan struct{}),
 			}
-		defer d.Halt()
-		if tt.halt > time.Nanosecond {
+
+			if tt.halt != 0 {
 				go func() {
-				time.Sleep(tt.halt)
-				d.Halt()
+					if err := d.Halt(); err != nil {
+						t.Errorf("Halt() expected nil but got %v", err)
+					}
 				}()
-		} else if tt.halt != 0 {
-			d.Halt()
 			}
 
-		got := d.writeVirtualRegister(0x04, 0xFF)
+			got := d.writeVirtualRegister(ctx, 0x04, 0xFF)
 			if _, ok := tt.wantErr.(*IOError); ok {
 				if _, ok := got.(*IOError); !ok {
 					t.Errorf("expected IOError but %T", got)
@@ -494,6 +480,7 @@ func TestDev_writeVirtualRegister(t *testing.T) {
 			} else if got != tt.wantErr {
 				t.Errorf("expected error: %v but got: %v", tt.wantErr, got)
 			}
+		})
 	}
 }
 
@@ -513,9 +500,14 @@ func TestDev_readVirtualRegister(t *testing.T) {
 		},
 		{
 			name: "errHalted",
-			tx:      []i2ctest.IO{},
+			tx: []i2ctest.IO{
+				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
+				{Addr: 0x49, W: []byte{writeReg, 0x04}, R: []byte{}},
+				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
+				{Addr: 0x49, W: []byte{readReg}, R: []byte{0x00}},
+			},
 			data:    []byte{0x00},
-			halt:    time.Nanosecond,
+			halt:    1,
 			timeout: time.Millisecond * 100,
 			wantErr: errHalted,
 		},
@@ -523,7 +515,6 @@ func TestDev_readVirtualRegister(t *testing.T) {
 			name: "errClearingBuffer",
 			tx: []i2ctest.IO{
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
-				// {Addr: 0x49, W: []byte{statusReg}, R: []byte{0x02}},
 			},
 			data:    []byte{0x00},
 			timeout: time.Millisecond * 100,
@@ -591,26 +582,36 @@ func TestDev_readVirtualRegister(t *testing.T) {
 		},
 	}
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			bus := &i2ctest.Playback{
 				Ops:       tt.tx,
 				DontPanic: true,
 			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+
+			sensorTimeout = tt.timeout
+			defer func() {
+				// cleanup
+				sensorTimeout = defaultSensorTimeOut
+				cancel()
+			}()
+
 			d := &Dev{
 				c:      &i2c.Dev{Bus: bus, Addr: 0x49},
-			done:    make(chan struct{}, 1),
-			timeout: tt.timeout,
+				cancel: cancel,
+				done:   make(chan struct{}),
 			}
-		// defer d.Halt()
-		if tt.halt > time.Nanosecond {
+
+			if tt.halt != 0 {
 				go func() {
-				time.Sleep(tt.halt)
-				d.Halt()
+					if err := d.Halt(); err != nil {
+						t.Errorf("Halt() expected nil but got %v", err)
+					}
 				}()
-		} else if tt.halt != 0 {
-			d.Halt()
 			}
 
-		got := d.readVirtualRegister(0x04, tt.data)
+			got := d.readVirtualRegister(ctx, 0x04, tt.data)
 			if _, ok := tt.wantErr.(*IOError); ok {
 				if _, ok := got.(*IOError); !ok {
 					t.Errorf("expected IOError but %T", got)
@@ -621,6 +622,7 @@ func TestDev_readVirtualRegister(t *testing.T) {
 			} else if got != tt.wantErr {
 				t.Errorf("expected error: %v but got: %v", tt.wantErr, got)
 			}
+		})
 	}
 }
 
@@ -634,9 +636,14 @@ func TestDev_pollDataReady(t *testing.T) {
 	}{
 		{
 			name: "errHalted",
-			tx:      []i2ctest.IO{},
-			halt:    time.Nanosecond,
-			timeout: time.Millisecond * 100,
+			tx: []i2ctest.IO{
+				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
+				{Addr: 0x49, W: []byte{writeReg, controlReg}, R: []byte{}},
+				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
+				{Addr: 0x49, W: []byte{readReg}, R: []byte{0x01}},
+			},
+			halt:    1,
+			timeout: time.Millisecond * 1000,
 			wantErr: errHalted,
 		},
 		{
@@ -644,7 +651,7 @@ func TestDev_pollDataReady(t *testing.T) {
 			tx: []i2ctest.IO{
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{}},
 			},
-			timeout: time.Millisecond * 100,
+			timeout: time.Millisecond * 1000,
 			wantErr: &IOError{"reading status register", nil},
 		},
 		{
@@ -653,7 +660,7 @@ func TestDev_pollDataReady(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
 				{Addr: 0x49, W: []byte{writeReg}, R: []byte{}},
 			},
-			timeout: time.Millisecond * 100,
+			timeout: time.Millisecond * 1000,
 			wantErr: &IOError{"setting virtual register", nil},
 		},
 		{
@@ -663,7 +670,7 @@ func TestDev_pollDataReady(t *testing.T) {
 				{Addr: 0x49, W: []byte{writeReg, controlReg}, R: []byte{}},
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x00}},
 			},
-			timeout: time.Millisecond * 100,
+			timeout: time.Millisecond * 1000,
 			wantErr: &IOError{"reading status register", nil},
 		},
 		{
@@ -674,7 +681,7 @@ func TestDev_pollDataReady(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
 				{Addr: 0x49, W: []byte{}, R: []byte{}},
 			},
-			timeout: time.Millisecond * 100,
+			timeout: time.Millisecond * 1000,
 			wantErr: &IOError{"reading virtual register", nil},
 		},
 		{
@@ -696,7 +703,7 @@ func TestDev_pollDataReady(t *testing.T) {
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x01}},
 				{Addr: 0x49, W: []byte{readReg}, R: []byte{0x03}},
 			},
-			timeout: time.Millisecond * 100,
+			timeout: time.Millisecond * 1000,
 			wantErr: nil,
 		},
 		{
@@ -732,27 +739,36 @@ func TestDev_pollDataReady(t *testing.T) {
 		},
 	}
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			bus := &i2ctest.Playback{
 				Ops:       tt.tx,
 				DontPanic: true,
 			}
 
+			ctx, cancel := context.WithCancel(context.Background())
+
+			sensorTimeout = tt.timeout
+			defer func() {
+				// cleanup
+				sensorTimeout = defaultSensorTimeOut
+				cancel()
+			}()
+
 			d := &Dev{
 				c:      &i2c.Dev{Bus: bus, Addr: 0x49},
-			done:    make(chan struct{}, 1),
-			timeout: tt.timeout,
+				cancel: cancel,
+				done:   make(chan struct{}),
 			}
-		defer d.Halt()
-		if tt.halt > time.Nanosecond {
+
+			if tt.halt != 0 {
 				go func() {
-				time.Sleep(tt.halt)
-				d.Halt()
+					if err := d.Halt(); err != nil {
+						t.Errorf("Halt() expected nil but got %v", err)
+					}
 				}()
-		} else if tt.halt != 0 {
-			d.Halt()
 			}
 
-		got := d.pollDataReady()
+			got := d.pollDataReady(ctx)
 			if _, ok := tt.wantErr.(*IOError); ok {
 				if _, ok := got.(*IOError); !ok {
 					t.Errorf("expected IOError but %T", got)
@@ -763,6 +779,7 @@ func TestDev_pollDataReady(t *testing.T) {
 			} else if got != tt.wantErr {
 				t.Errorf("expected error: %v but got: %v", tt.wantErr, got)
 			}
+		})
 	}
 }
 
@@ -796,6 +813,12 @@ func TestNew(t *testing.T) {
 		if tt.want2 != d.interrupt {
 			t.Errorf("New() wanted %v but got %v", tt.want2, d.interrupt)
 		}
+
+		// Halt with empty context.
+		err = d.Halt()
+		if err != nil {
+			t.Errorf("New Sensor halt wanted nil but got %v", err)
+		}
 	}
 }
 
@@ -812,24 +835,10 @@ func TestDev_Gain(t *testing.T) {
 			name: "errStatusIO",
 			tx: []i2ctest.IO{
 				{Addr: 0x49, W: []byte{statusReg}, R: []byte{}},
-				// {Addr: 0x49, W: []byte{writeReg}, R: []byte{}},
 			},
 			timeout: time.Millisecond * 100,
 			wantErr: &IOError{"reading status register", nil},
 		},
-		{
-			name: "errHalt",
-			tx: []i2ctest.IO{
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x03}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x03}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x03}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x03}},
-				{Addr: 0x49, W: []byte{statusReg}, R: []byte{0x03}},
-			},
-			halt:    true,
-			timeout: time.Millisecond * 100,
-			wantErr: errHalted,
-		},
 		{
 			name:    "errGainValue",
 			gain:    Gain(255),
@@ -854,21 +863,19 @@ func TestDev_Gain(t *testing.T) {
 			DontPanic: true,
 		}
 		d, _ := New(bus, &DefaultOpts)
-		go func() {
-			time.Sleep(time.Millisecond * 1)
-			d.Halt()
-		}()
+
+		sensorTimeout = tt.timeout
 		got := d.Gain(tt.gain)
 
 		if _, ok := tt.wantErr.(*IOError); ok {
 			if _, ok := got.(*IOError); !ok {
-				t.Errorf("expected IOError but %T", got)
+				t.Errorf("%s: expected IOError but %T", tt.name, got)
 			}
 			if got.(*IOError).Op != tt.wantErr.(*IOError).Op {
-				t.Errorf("expected %s, but got %s", tt.wantErr.(*IOError).Op, got.(*IOError).Op)
+				t.Errorf("%s: expected %s, but got %s", tt.name, tt.wantErr.(*IOError).Op, got.(*IOError).Op)
 			}
 		} else if got != tt.wantErr {
-			t.Errorf("expected error: %v but got: %v", tt.wantErr, got)
+			t.Errorf("%s: expected error: %v but got: %v", tt.name, tt.wantErr, got)
 		}
 	}
 }
@@ -940,20 +947,3 @@ func TestIOError_Error(t *testing.T) {
 		}
 	}
 }
-
-func TestIntergration_AfterHalt(t *testing.T) {
-	bus := &i2ctest.Playback{
-		Ops:       sensorTestCaseValidRead,
-		DontPanic: true,
-	}
-	d, _ := New(bus, &DefaultOpts)
-
-	d.Halt()
-	if _, err := d.Sense(physic.MilliAmpere*100, time.Millisecond*3); err == errHalted {
-		t.Errorf("got %v expected nil", errHalted)
-	}
-	d.Halt()
-	if err := d.Gain(G1x); err == errHalted {
-		t.Errorf("got %v expected nil", errHalted)
-	}
-}