[DATA-3338] - collector test improvements (#4551)

viamrobotics · Nov 14, 2024 · 1460f8c · 1460f8c
1 parent 8bbc629
commit 1460f8c
Show file tree

Hide file tree

Showing 27 changed files with 853 additions and 363 deletions.
diff --git a/components/arm/collector.go → components/arm/collectors.go b/components/arm/collector.go → components/arm/collectors.go
diff --git a/components/arm/collectors_test.go b/components/arm/collectors_test.go
@@ -5,9 +5,9 @@ import (
 	"testing"
 	"time"
 
-	clk "github.com/benbjohnson/clock"
+	"github.com/benbjohnson/clock"
 	"github.com/golang/geo/r3"
-	v1 "go.viam.com/api/common/v1"
+	datasyncpb "go.viam.com/api/app/datasync/v1"
 	pb "go.viam.com/api/component/arm/v1"
 	"go.viam.com/test"
 
@@ -22,8 +22,7 @@ import (
 
 const (
 	componentName   = "arm"
-	captureInterval = time.Second
-	numRetries      = 5
+	captureInterval = time.Millisecond
 )
 
 var floatList = &pb.JointPositions{Values: []float64{1.0, 2.0, 3.0}}
@@ -32,40 +31,50 @@ func TestCollectors(t *testing.T) {
 	tests := []struct {
 		name      string
 		collector data.CollectorConstructor
-		expected  map[string]any
+		expected  *datasyncpb.SensorData
 	}{
 		{
 			name:      "End position collector should write a pose",
 			collector: arm.NewEndPositionCollector,
-			expected: tu.ToProtoMapIgnoreOmitEmpty(pb.GetEndPositionResponse{
-				Pose: &v1.Pose{
-					OX:    0,
-					OY:    0,
-					OZ:    1,
-					Theta: 0,
-					X:     1,
-					Y:     2,
-					Z:     3,
-				},
-			}),
+			expected: &datasyncpb.SensorData{
+				Metadata: &datasyncpb.SensorMetadata{},
+				Data: &datasyncpb.SensorData_Struct{Struct: tu.ToStructPBStruct(t, map[string]any{
+					"pose": map[string]any{
+						"o_x":   0,
+						"o_y":   0,
+						"o_z":   1,
+						"theta": 0,
+						"x":     1,
+						"y":     2,
+						"z":     3,
+					},
+				})},
+			},
 		},
 		{
 			name:      "Joint positions collector should write a list of positions",
 			collector: arm.NewJointPositionsCollector,
-			expected:  tu.ToProtoMapIgnoreOmitEmpty(pb.GetJointPositionsResponse{Positions: floatList}),
+			expected: &datasyncpb.SensorData{
+				Metadata: &datasyncpb.SensorMetadata{},
+				Data: &datasyncpb.SensorData_Struct{Struct: tu.ToStructPBStruct(t, map[string]any{
+					"positions": map[string]any{
+						"values": []any{1.0, 2.0, 3.0},
+					},
+				})},
+			},
 		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			mockClock := clk.NewMock()
-			buf := tu.MockBuffer{}
+			start := time.Now()
+			buf := tu.NewMockBuffer()
 			params := data.CollectorParams{
 				ComponentName: componentName,
 				Interval:      captureInterval,
 				Logger:        logging.NewTestLogger(t),
-				Clock:         mockClock,
-				Target:        &buf,
+				Clock:         clock.New(),
+				Target:        buf,
 			}
 
 			arm := newArm()
@@ -74,13 +83,11 @@ func TestCollectors(t *testing.T) {
 
 			defer col.Close()
 			col.Collect()
-			mockClock.Add(captureInterval)
 
-			tu.Retry(func() bool {
-				return buf.Length() != 0
-			}, numRetries)
-			test.That(t, buf.Length(), test.ShouldBeGreaterThan, 0)
-			test.That(t, buf.Writes[0].GetStruct().AsMap(), test.ShouldResemble, tc.expected)
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			tu.CheckMockBufferWrites(t, ctx, start, buf.Writes, tc.expected)
+			buf.Close()
 		})
 	}
 }

diff --git a/components/board/collector.go → components/board/collectors.go b/components/board/collector.go → components/board/collectors.go
diff --git a/components/board/collectors_test.go b/components/board/collectors_test.go
@@ -6,9 +6,9 @@ import (
 	"testing"
 	"time"
 
-	clk "github.com/benbjohnson/clock"
+	"github.com/benbjohnson/clock"
 	"github.com/golang/protobuf/ptypes/wrappers"
-	pb "go.viam.com/api/component/board/v1"
+	datasyncpb "go.viam.com/api/app/datasync/v1"
 	"go.viam.com/test"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
@@ -22,18 +22,15 @@ import (
 
 const (
 	componentName   = "board"
-	captureInterval = time.Second
-	numRetries      = 5
+	captureInterval = time.Millisecond
 )
 
 func TestCollectors(t *testing.T) {
 	tests := []struct {
-		name          string
-		params        data.CollectorParams
-		collector     data.CollectorConstructor
-		expected      map[string]any
-		shouldError   bool
-		expectedError error
+		name      string
+		params    data.CollectorParams
+		collector data.CollectorConstructor
+		expected  *datasyncpb.SensorData
 	}{
 		{
 			name: "Board analog collector should write an analog response",
@@ -46,13 +43,15 @@ func TestCollectors(t *testing.T) {
 				},
 			},
 			collector: board.NewAnalogCollector,
-			expected: tu.ToProtoMapIgnoreOmitEmpty(pb.ReadAnalogReaderResponse{
-				Value:    1,
-				MinRange: 0,
-				MaxRange: 10,
-				StepSize: 0.1,
-			}),
-			shouldError: false,
+			expected: &datasyncpb.SensorData{
+				Metadata: &datasyncpb.SensorMetadata{},
+				Data: &datasyncpb.SensorData_Struct{Struct: tu.ToStructPBStruct(t, map[string]any{
+					"value":     1,
+					"min_range": 0,
+					"max_range": 10,
+					"step_size": float64(float32(0.1)),
+				})},
+			},
 		},
 		{
 			name: "Board gpio collector should write a gpio response",
@@ -65,33 +64,33 @@ func TestCollectors(t *testing.T) {
 				},
 			},
 			collector: board.NewGPIOCollector,
-			expected: tu.ToProtoMapIgnoreOmitEmpty(pb.GetGPIOResponse{
-				High: true,
-			}),
-			shouldError: false,
+			expected: &datasyncpb.SensorData{
+				Metadata: &datasyncpb.SensorMetadata{},
+				Data: &datasyncpb.SensorData_Struct{Struct: tu.ToStructPBStruct(t, map[string]any{
+					"high": true,
+				})},
+			},
 		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			mockClock := clk.NewMock()
-			buf := tu.MockBuffer{}
-			tc.params.Clock = mockClock
-			tc.params.Target = &buf
+			start := time.Now()
+			buf := tu.NewMockBuffer()
+			tc.params.Clock = clock.New()
+			tc.params.Target = buf
 
 			board := newBoard()
 			col, err := tc.collector(board, tc.params)
 			test.That(t, err, test.ShouldBeNil)
 
 			defer col.Close()
 			col.Collect()
-			mockClock.Add(captureInterval)
 
-			tu.Retry(func() bool {
-				return buf.Length() != 0
-			}, numRetries)
-			test.That(t, buf.Length(), test.ShouldBeGreaterThan, 0)
-			test.That(t, buf.Writes[0].GetStruct().AsMap(), test.ShouldResemble, tc.expected)
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			tu.CheckMockBufferWrites(t, ctx, start, buf.Writes, tc.expected)
+			buf.Close()
 		})
 	}
 }

diff --git a/components/camera/collector.go → components/camera/collectors.go b/components/camera/collector.go → components/camera/collectors.go
@@ -36,8 +36,6 @@ func (m method) String() string {
 	return "Unknown"
 }
 
-// TODO: add tests for this file.
-
 func newNextPointCloudCollector(resource interface{}, params data.CollectorParams) (data.Collector, error) {
 	camera, err := assertCamera(resource)
 	if err != nil {