-
Notifications
You must be signed in to change notification settings - Fork 3.7k
/
datadriven_test.go
138 lines (122 loc) · 4.64 KB
/
datadriven_test.go
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
// Copyright 2022 The Cockroach Authors.
//
// Licensed as a CockroachDB Enterprise file under the Cockroach Community
// License (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
package spanconfigsplitterccl
import (
"context"
"fmt"
"strings"
"testing"
"github.com/cockroachdb/cockroach/pkg/base"
_ "github.com/cockroachdb/cockroach/pkg/ccl/kvccl/kvtenantccl"
_ "github.com/cockroachdb/cockroach/pkg/ccl/partitionccl"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/spanconfig"
"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigsplitter"
"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigtestutils/spanconfigtestcluster"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/datadriven"
"github.com/stretchr/testify/require"
)
// TestDataDriven is a data-driven test for spanconfig.Splitter. It offers
// the following commands:
//
// - "exec-sql"
// Executes the input SQL query.
//
// - "query-sql"
// Executes the input SQL query and prints the results.
//
// - "splits" [database=<str> table=<str>] [id=<int>]
// Prints the number splits generated the referenced object (named database +
// table, or descriptor id). Also logs the set of internal steps the Splitter
// takes to arrive at the number.
func TestDataDriven(t *testing.T) {
defer leaktest.AfterTest(t)()
ctx := context.Background()
var steps strings.Builder
scKnobs := &spanconfig.TestingKnobs{
// Instead of relying on the GC job to wait out TTLs and clear out descriptors,
// let's simply exclude dropped tables to simulate descriptors no longer existing.
// See comment on ExcludeDroppedDescriptorsFromLookup for more details.
ExcludeDroppedDescriptorsFromLookup: true,
// We run the reconciler manually in this test (through the span config
// test cluster).
ManagerDisableJobCreation: true,
// SplitterStepLogger captures splitter-internal steps for test output.
SplitterStepLogger: func(step string) {
steps.WriteString(fmt.Sprintf("%s\n", step))
},
}
datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) {
defer log.Scope(t).Close(t)
tc := testcluster.StartTestCluster(t, 1, base.TestClusterArgs{
ServerArgs: base.TestServerArgs{
DefaultTestTenant: base.TestControlsTenantsExplicitly,
Knobs: base.TestingKnobs{
SpanConfig: scKnobs,
},
},
})
defer tc.Stopper().Stop(ctx)
spanConfigTestCluster := spanconfigtestcluster.NewHandle(t, tc, scKnobs)
defer spanConfigTestCluster.Cleanup()
tenantID := roachpb.MustMakeTenantID(10)
tenant := spanConfigTestCluster.InitializeTenant(ctx, tenantID)
// TODO(irfansharif): Expose this through the test harness once we integrate
// it into the schema changer.
splitter := spanconfigsplitter.New(tenant.ExecCfg().Codec, scKnobs)
var lastSeenID descpb.ID
datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
switch d.Cmd {
case "exec-sql":
tenant.Exec(d.Input)
case "query-sql":
rows := tenant.Query(d.Input)
output, err := sqlutils.RowsToDataDrivenOutput(rows)
require.NoError(t, err)
return output
case "splits":
// Parse the args to get the object ID we're looking to split.
var objID descpb.ID
switch {
case d.HasArg("id"):
var scanID int
d.ScanArgs(t, "id", &scanID)
objID = descpb.ID(scanID)
case d.HasArg("database"):
// NB: Name resolution for dropped descriptors does not work like
// you'd expect, i.e. it does not work at all. We have to look things
// up by ID first.
var dbName, tbName string
d.ScanArgs(t, "database", &dbName)
d.ScanArgs(t, "table", &tbName)
tbl := tenant.LookupTableByName(ctx, dbName, tbName)
objID = tbl.GetID()
case d.HasArg("last_seen_id"):
objID = lastSeenID
default:
d.Fatalf(t, "insufficient/improper args (%v) provided to split", d.CmdArgs)
}
steps.Reset()
lastSeenID = objID
splits, err := splitter.Splits(ctx, tenant.LookupTableDescriptorByID(ctx, objID))
require.NoError(t, err)
steps.WriteString(fmt.Sprintf("= %d", splits))
return steps.String()
default:
t.Fatalf("unknown command: %s", d.Cmd)
}
return ""
})
})
}