plugin-sdk/plugin/nulls_test.go at main · cloudquery/plugin-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package plugin

import (
	"testing"
	"time"

	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/cloudquery/plugin-sdk/v4/schema"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestWithTestIgnoreNullsInLists(t *testing.T) {
	s := &WriterTestSuite{ignoreNullsInLists: true}

	tg := schema.NewTestDataGenerator(0)
	source := schema.TestTable("ignore_nulls_in_lists", schema.TestSourceOptions{})
	resource := s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
		SourceName: "allow_null",
		SyncTime:   time.Now(),
		MaxRows:    100,
		NullRows:   false,
	}))
	for _, c := range resource.Columns() {
		assertNoNullsInLists(t, c)
	}

	resource = s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
		SourceName: "ignore_nulls_in_lists",
		SyncTime:   time.Now(),
		MaxRows:    100,
		NullRows:   true,
	}))
	for _, c := range resource.Columns() {
		assertNoNullsInLists(t, c)
	}
}

func assertNoNullsInLists(t *testing.T, arr arrow.Array) {
	// traverse
	switch arr := arr.(type) {
	case array.ListLike:
		assert.Zero(t, arr.ListValues().NullN())
	case *array.Struct:
		for i := 0; i < arr.NumField(); i++ {
			assertNoNullsInLists(t, arr.Field(i))
		}
	}
}

func TestWithTestSourceAllowNull(t *testing.T) {
	s := &WriterTestSuite{allowNull: func(dt arrow.DataType) bool {
		switch dt.(type) {
		case *arrow.StructType, arrow.ListLikeType:
			return false
		default:
			return true
		}
	}}

	tg := schema.NewTestDataGenerator(0)
	source := schema.TestTable("allow_null", schema.TestSourceOptions{})
	resource := s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
		SourceName: "allow_null",
		SyncTime:   time.Now(),
		MaxRows:    100,
		NullRows:   false,
	}))
	for _, c := range resource.Columns() {
		assertNoNulls(t, s.allowNull, c)
	}

	resource = s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
		SourceName: "allow_null",
		SyncTime:   time.Now(),
		MaxRows:    100,
		NullRows:   true,
	}))
	for _, c := range resource.Columns() {
		assertNoNulls(t, s.allowNull, c)
	}
}

func assertNoNulls(t *testing.T, allowNull AllowNullFunc, arr arrow.Array) {
	require.NotNil(t, allowNull)

	if !allowNull(arr.DataType()) {
		assert.Zero(t, arr.NullN())
	}

	// traverse
	switch arr := arr.(type) {
	case array.ListLike:
		assertNoNulls(t, allowNull, arr.ListValues())
	case *array.Struct:
		for i := 0; i < arr.NumField(); i++ {
			assertNoNulls(t, allowNull, arr.Field(i))
		}
	}
}