@@ -13,6 +13,7 @@ import (
1313 "github.com/cloudquery/plugin-sdk/v4/plugin"
1414 "github.com/cloudquery/plugin-sdk/v4/schema"
1515 "github.com/google/uuid"
16+ "github.com/stretchr/testify/require"
1617)
1718
1819func TestWriteRead (t * testing.T ) {
@@ -40,7 +41,7 @@ func TestWriteRead(t *testing.T) {
4041 StableTime : time .Date (2021 , 1 , 2 , 0 , 0 , 0 , 0 , time .UTC ),
4142 }
4243 tg := schema .NewTestDataGenerator ()
43- records := tg .Generate (table , opts )
44+ record := tg .Generate (table , opts )
4445
4546 cl , err := NewClient (tc .options ... )
4647 if err != nil {
@@ -51,7 +52,7 @@ func TestWriteRead(t *testing.T) {
5152 writer := bufio .NewWriter (& b )
5253 reader := bufio .NewReader (& b )
5354
54- if err := types .WriteAll (cl , writer , table , records ); err != nil {
55+ if err := types .WriteAll (cl , writer , table , []arrow. Record { record } ); err != nil {
5556 t .Fatal (err )
5657 }
5758 writer .Flush ()
@@ -74,19 +75,13 @@ func TestWriteRead(t *testing.T) {
7475 readErr = cl .Read (byteReader , table , ch )
7576 close (ch )
7677 }()
77- totalCount := 0
78+ received := make ([]arrow. Record , 0 , tc . outputCount )
7879 for got := range ch {
79- if diff := plugin .RecordDiff (records [totalCount ], got ); diff != "" {
80- t .Errorf ("got diff: %s" , diff )
81- }
82- totalCount ++
83- }
84- if readErr != nil {
85- t .Fatal (readErr )
86- }
87- if totalCount != tc .outputCount {
88- t .Errorf ("got %d row(s), want %d" , totalCount , tc .outputCount )
80+ received = append (received , got )
8981 }
82+ require .Empty (t , plugin .RecordsDiff (table .ToArrowSchema (), []arrow.Record {record }, received ))
83+ require .NoError (t , readErr )
84+ require .Equalf (t , tc .outputCount , len (received ), "got %d row(s), want %d" , len (received ), tc .outputCount )
9085 })
9186 }
9287}
@@ -101,7 +96,7 @@ func BenchmarkWrite(b *testing.B) {
10196 MaxRows : 1000 ,
10297 }
10398 tg := schema .NewTestDataGenerator ()
104- records := tg .Generate (table , opts )
99+ record := tg .Generate (table , opts )
105100
106101 cl , err := NewClient ()
107102 if err != nil {
@@ -111,7 +106,7 @@ func BenchmarkWrite(b *testing.B) {
111106 writer := bufio .NewWriter (& buf )
112107 b .ResetTimer ()
113108 for i := 0 ; i < b .N ; i ++ {
114- if err := types .WriteAll (cl , writer , table , records ); err != nil {
109+ if err := types .WriteAll (cl , writer , table , []arrow. Record { record } ); err != nil {
115110 b .Fatal (err )
116111 }
117112
0 commit comments