github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/workloadsql/workloadsql_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package workloadsql
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"strconv"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/base"
    20  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    21  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    22  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    23  	"github.com/cockroachdb/cockroach/pkg/util/uuid"
    24  	"github.com/cockroachdb/cockroach/pkg/workload"
    25  	"github.com/cockroachdb/cockroach/pkg/workload/bank"
    26  )
    27  
    28  func TestSetup(t *testing.T) {
    29  	defer leaktest.AfterTest(t)()
    30  
    31  	tests := []struct {
    32  		rows        int
    33  		batchSize   int
    34  		concurrency int
    35  	}{
    36  		{10, 1, 1},
    37  		{10, 9, 1},
    38  		{10, 10, 1},
    39  		{10, 100, 1},
    40  		{10, 1, 4},
    41  		{10, 9, 4},
    42  		{10, 10, 4},
    43  		{10, 100, 4},
    44  	}
    45  
    46  	ctx := context.Background()
    47  	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{UseDatabase: `test`})
    48  	defer s.Stopper().Stop(ctx)
    49  	sqlutils.MakeSQLRunner(db).Exec(t, `CREATE DATABASE test`)
    50  
    51  	for _, test := range tests {
    52  		t.Run(fmt.Sprintf("rows=%d/batch=%d", test.rows, test.batchSize), func(t *testing.T) {
    53  			sqlDB := sqlutils.MakeSQLRunner(db)
    54  			sqlDB.Exec(t, `DROP TABLE IF EXISTS bank`)
    55  
    56  			gen := bank.FromRows(test.rows)
    57  			l := InsertsDataLoader{BatchSize: test.batchSize, Concurrency: test.concurrency}
    58  			if _, err := Setup(ctx, db, gen, l); err != nil {
    59  				t.Fatalf("%+v", err)
    60  			}
    61  
    62  			for _, table := range gen.Tables() {
    63  				var c int
    64  				sqlDB.QueryRow(t, fmt.Sprintf(`SELECT count(*) FROM %s`, table.Name)).Scan(&c)
    65  				// There happens to be 1 row per batch in bank.
    66  				if c != table.InitialRows.NumBatches {
    67  					t.Errorf(`%s: got %d rows expected %d`,
    68  						table.Name, c, table.InitialRows.NumBatches)
    69  				}
    70  			}
    71  		})
    72  	}
    73  }
    74  
    75  func TestSplits(t *testing.T) {
    76  	defer leaktest.AfterTest(t)()
    77  
    78  	ctx := context.Background()
    79  	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{UseDatabase: `test`})
    80  	defer s.Stopper().Stop(ctx)
    81  	sqlutils.MakeSQLRunner(db).Exec(t, `CREATE DATABASE test`)
    82  
    83  	for _, ranges := range []int{1, 2, 3, 4, 10} {
    84  
    85  		tables := []workload.Table{
    86  			{
    87  				Name:   `ints`,
    88  				Schema: `(a INT PRIMARY KEY)`,
    89  				Splits: workload.Tuples(ranges-1, func(i int) []interface{} {
    90  					return []interface{}{i}
    91  				}),
    92  			},
    93  			{
    94  				Name:   `floats`,
    95  				Schema: `(a FLOAT PRIMARY KEY)`,
    96  				Splits: workload.Tuples(ranges-1, func(i int) []interface{} {
    97  					return []interface{}{float64(i)}
    98  				}),
    99  			},
   100  			{
   101  				Name:   `strings`,
   102  				Schema: `(a STRING PRIMARY KEY)`,
   103  				Splits: workload.Tuples(ranges-1, func(i int) []interface{} {
   104  					return []interface{}{strconv.Itoa(i)}
   105  				}),
   106  			},
   107  			{
   108  				Name:   `bytes`,
   109  				Schema: `(a BYTES PRIMARY KEY)`,
   110  				Splits: workload.Tuples(ranges-1, func(i int) []interface{} {
   111  					return []interface{}{strconv.Itoa(i)}
   112  				}),
   113  			},
   114  			{
   115  				Name:   `uuids`,
   116  				Schema: `(a UUID PRIMARY KEY)`,
   117  				Splits: workload.Tuples(ranges-1, func(i int) []interface{} {
   118  					u, err := uuid.NewV4()
   119  					if err != nil {
   120  						panic(err)
   121  					}
   122  					return []interface{}{u.String()}
   123  				}),
   124  			},
   125  		}
   126  
   127  		t.Run(fmt.Sprintf("ranges=%d", ranges), func(t *testing.T) {
   128  			sqlDB := sqlutils.MakeSQLRunner(db)
   129  			for _, table := range tables {
   130  				sqlDB.Exec(t, fmt.Sprintf(`DROP TABLE IF EXISTS %s`, table.Name))
   131  				sqlDB.Exec(t, fmt.Sprintf(`CREATE TABLE %s %s`, table.Name, table.Schema))
   132  
   133  				const concurrency = 10
   134  				if err := Split(ctx, db, table, concurrency); err != nil {
   135  					t.Fatalf("%+v", err)
   136  				}
   137  
   138  				countRangesQ := fmt.Sprintf(
   139  					`SELECT count(*) FROM [SHOW RANGES FROM TABLE test.%s]`, table.Name,
   140  				)
   141  				var actual int
   142  				sqlDB.QueryRow(t, countRangesQ).Scan(&actual)
   143  				if ranges != actual {
   144  					t.Errorf(`expected %d got %d`, ranges, actual)
   145  				}
   146  			}
   147  		})
   148  	}
   149  }