add for both methods

Author: maaarcelino

plugin/testing_write_delete.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/samber/lo"
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/array"
 	"github.com/apache/arrow-go/v18/arrow/memory"
@@ -177,55 +178,58 @@ func (s *WriterTestSuite) testDeleteRecordBasic(ctx context.Context, t *testing.
 	r.NoErrorf(err, "failed to read")
 	r.EqualValuesf(1, TotalRows(records), "unexpected amount of items")
 
-	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, 1)), "failed to delete record")
+	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, toDeleteValue(table, 1))), "failed to delete record")
 
 	records, err = s.plugin.readAll(ctx, table)
 	r.NoErrorf(err, "failed to read after delete with no match")
 	r.EqualValuesf(1, TotalRows(records), "unexpected amount of items after delete with no match")
 	r.Emptyf(RecordsDiff(table.ToArrowSchema(), records, []arrow.Record{record1}), "record differs after delete with no match")
 
-	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, 0)), "failed to delete record")
+	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, toDeleteValue(table, 0))), "failed to delete record")
 
 	records, err = s.plugin.readAll(ctx, table)
 	r.NoErrorf(err, "failed to read after delete with match")
 	r.EqualValuesf(0, TotalRows(records), "unexpected amount of items after delete with match")
 }
 
-func (s *WriterTestSuite) toMessageBatch(table *schema.Table, val int64) []message.WriteMessage {
-	batchSize := 1
-	if s.tests.DeleteRecordBatchSize > 0 {
-		batchSize = s.tests.DeleteRecordBatchSize
-	}
-
-	messages := make([]message.WriteMessage, batchSize)
-
+func toDeleteValue(table *schema.Table, val int64) *arrow.Record {
 	builder := array.NewRecordBuilder(memory.DefaultAllocator, (&schema.Table{
 		Name: table.Name,
 		Columns: schema.ColumnList{
 			schema.Column{Name: "id", Type: arrow.PrimitiveTypes.Int64},
 		},
 	}).ToArrowSchema())
 	builder.Field(0).(*array.Int64Builder).Append(val)
-	deleteValue := builder.NewRecord()
+	return lo.ToPtr(builder.NewRecord())
+}
+
+func (s *WriterTestSuite) toMessageBatch(table *schema.Table, deleteValue *arrow.Record) []message.WriteMessage {
+	batchSize := 1
+	if s.tests.DeleteRecordBatchSize > 0 {
+		batchSize = s.tests.DeleteRecordBatchSize
+	}
+
+	messages := make([]message.WriteMessage, batchSize)
 
 	for i, _ := range messages {
-		messages[i] = &message.WriteDeleteRecord{
-			DeleteRecord: message.DeleteRecord{
-				TableName: table.Name,
-				WhereClause: message.PredicateGroups{
-					{
-						GroupingType: "AND",
-						Predicates: []message.Predicate{
-							{
-								Operator: "eq",
-								Column:   "id",
-								Record:   deleteValue,
-							},
+		record := message.DeleteRecord{
+			TableName: table.Name,
+		}
+		if deleteValue != nil {
+			record.WhereClause = message.PredicateGroups{
+				{
+					GroupingType: "AND",
+					Predicates: []message.Predicate{
+						{
+							Operator: "eq",
+							Column:   "id",
+							Record:   *deleteValue,
 						},
 					},
 				},
-			},
+			}
 		}
+		messages[i] = &message.WriteDeleteRecord{DeleteRecord: record}
 	}
 
 	return messages
@@ -259,11 +263,7 @@ func (s *WriterTestSuite) testDeleteAllRecords(ctx context.Context, t *testing.T
 	r.NoErrorf(err, "failed to read")
 	r.EqualValuesf(1, TotalRows(records), "unexpected amount of items")
 
-	r.NoErrorf(s.plugin.writeOne(ctx, &message.WriteDeleteRecord{
-		DeleteRecord: message.DeleteRecord{
-			TableName: table.Name,
-		},
-	}), "failed to delete records")
+	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, nil)), "failed to delete record")
 
 	records, err = s.plugin.readAll(ctx, table)
 	r.NoErrorf(err, "failed to read after delete all records")
@@ -276,11 +276,7 @@ func (s *WriterTestSuite) testDeleteAllRecords(ctx context.Context, t *testing.T
 
 	r.NoErrorf(s.plugin.writeOne(ctx, &message.WriteInsert{Record: record2}), "failed to insert second record")
 
-	r.NoErrorf(s.plugin.writeOne(ctx, &message.WriteDeleteRecord{
-		DeleteRecord: message.DeleteRecord{
-			TableName: table.Name,
-		},
-	}), "failed to delete records second time")
+	r.NoError(s.plugin.WriteAll(ctx, s.toMessageBatch(table, nil)), "failed to delete record")
 
 	records, err = s.plugin.readAll(ctx, table)
 	r.NoErrorf(err, "failed to read second time")