kv: move test_utils file to new kvtestutils package

Release note: None
cockroachdb · Feb 14, 2025 · 8cf8395 · 8cf8395
1 parent b374872
commit 8cf8395
Show file tree

Hide file tree

Showing 13 changed files with 34 additions and 17 deletions.
diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel
@@ -1586,6 +1586,7 @@ GO_TARGETS = [
     "//pkg/kv/kvserver/uncertainty:uncertainty_test",
     "//pkg/kv/kvserver:kvserver",
     "//pkg/kv/kvserver:kvserver_test",
+    "//pkg/kv/kvtestutils:kvtestutils",
     "//pkg/kv:kv",
     "//pkg/kv:kv_test",
     "//pkg/multitenant/mtinfo:mtinfo",

diff --git a/pkg/ccl/kvccl/kvfollowerreadsccl/BUILD.bazel b/pkg/ccl/kvccl/kvfollowerreadsccl/BUILD.bazel
@@ -59,6 +59,7 @@ go_test(
         "//pkg/kv/kvserver/closedts",
         "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/kv/kvserver/kvserverbase",
+        "//pkg/kv/kvtestutils",
         "//pkg/roachpb",
         "//pkg/rpc",
         "//pkg/security/securityassets",

diff --git a/pkg/ccl/kvccl/kvfollowerreadsccl/boundedstaleness_test.go b/pkg/ccl/kvccl/kvfollowerreadsccl/boundedstaleness_test.go
@@ -16,9 +16,9 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/ccl"
-	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvbase"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -248,7 +248,7 @@ func (bse *boundedStalenessEvents) onStmtTrace(nodeIdx int, rec tracingpb.Record
 					operation:    spans[sp.ParentSpanID].Operation,
 					nodeIdx:      nodeIdx,
 					localRead:    tracing.LogsContainMsg(sp, kvbase.RoutingRequestLocallyMsg),
-					followerRead: kv.OnlyFollowerReads(rec),
+					followerRead: kvtestutils.OnlyFollowerReads(rec),
 					remoteLeaseholderRead: tracing.LogsContainMsg(sp, "[NotLeaseHolderError] lease held by different store;") &&
 						tracing.LogsContainMsg(sp, "trying next peer"),
 				})

diff --git a/pkg/ccl/kvccl/kvfollowerreadsccl/followerreads_test.go b/pkg/ccl/kvccl/kvfollowerreadsccl/followerreads_test.go
@@ -26,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/server"
@@ -890,7 +891,7 @@ func TestFollowerReadsWithStaleDescriptor(t *testing.T) {
 	n4.Exec(t, historicalQuery.Load().(string))
 	// As a sanity check, verify that this was not a follower read.
 	rec := <-recCh
-	require.False(t, kv.OnlyFollowerReads(rec), "query was served through follower reads: %s", rec)
+	require.False(t, kvtestutils.OnlyFollowerReads(rec), "query was served through follower reads: %s", rec)
 	// Check that the cache was properly updated.
 	entry, err = n4Cache.TestingGetCached(ctx, tablePrefix, false, roachpb.LAG_BY_CLUSTER_SETTING)
 	require.NoError(t, err)
@@ -917,7 +918,7 @@ func TestFollowerReadsWithStaleDescriptor(t *testing.T) {
 	rec = <-recCh
 
 	// Look at the trace and check that we've served a follower read.
-	require.True(t, kv.OnlyFollowerReads(rec), "query was not served through follower reads: %s", rec)
+	require.True(t, kvtestutils.OnlyFollowerReads(rec), "query was not served through follower reads: %s", rec)
 
 	// Check that the follower read metric was incremented.
 	var followerReadsCountAfter int64
@@ -963,7 +964,7 @@ func TestFollowerReadsWithStaleDescriptor(t *testing.T) {
 	// Sanity check that the plan was distributed.
 	require.True(t, strings.Contains(rec.String(), "creating DistSQL plan with isLocal=false"))
 	// Look at the trace and check that we've served a follower read.
-	require.True(t, kv.OnlyFollowerReads(rec), "query was not served through follower reads: %s", rec)
+	require.True(t, kvtestutils.OnlyFollowerReads(rec), "query was not served through follower reads: %s", rec)
 	// Verify that we didn't produce the "misplanned ranges" metadata that would
 	// purge the non-stale entries from the range cache on n4.
 	require.False(t, strings.Contains(rec.String(), "clearing entries overlapping"))

diff --git a/pkg/cmd/roachtest/tests/BUILD.bazel b/pkg/cmd/roachtest/tests/BUILD.bazel
@@ -242,8 +242,8 @@ go_library(
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/keys",
-        "//pkg/kv",
         "//pkg/kv/kvpb",
+        "//pkg/kv/kvtestutils",
         "//pkg/multitenant/mtinfopb",
         "//pkg/roachpb",
         "//pkg/roachprod",

diff --git a/pkg/cmd/roachtest/tests/kv.go b/pkg/cmd/roachtest/tests/kv.go
@@ -21,7 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
-	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -924,7 +924,7 @@ func registerKVRangeLookups(r registry.Registry) {
 							EXPERIMENTAL_RELOCATE
 								SELECT ARRAY[$1, $2, $3], CAST(floor(random() * 9223372036854775808) AS INT)
 						`, newReplicas[0]+1, newReplicas[1]+1, newReplicas[2]+1)
-						if err != nil && !pgerror.IsSQLRetryableError(err) && !kv.IsExpectedRelocateError(err) {
+						if err != nil && !pgerror.IsSQLRetryableError(err) && !kvtestutils.IsExpectedRelocateError(err) {
 							return err
 						}
 					default:

diff --git a/pkg/kv/BUILD.bazel b/pkg/kv/BUILD.bazel
@@ -10,7 +10,6 @@ go_library(
         "mock_transactional_sender.go",
         "range_lookup.go",
         "sender.go",
-        "test_utils.go",
         "txn.go",
         "util.go",
     ],
@@ -19,7 +18,6 @@ go_library(
     deps = [
         "//pkg/base",
         "//pkg/keys",
-        "//pkg/kv/kvbase",
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/closedts",
         "//pkg/kv/kvserver/concurrency/isolation",
@@ -28,7 +26,6 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/sql/sessiondatapb",
         "//pkg/storage/enginepb",
-        "//pkg/testutils",
         "//pkg/util/admission",
         "//pkg/util/admission/admissionpb",
         "//pkg/util/duration",
@@ -40,8 +37,6 @@ go_library(
         "//pkg/util/stop",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
-        "//pkg/util/tracing",
-        "//pkg/util/tracing/tracingpb",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_apd_v3//:apd",
         "@com_github_cockroachdb_errors//:errors",
@@ -75,6 +70,7 @@ go_test(
         "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/kv/kvserver/kvserverbase",
         "//pkg/kv/kvserver/liveness/livenesspb",
+        "//pkg/kv/kvtestutils",
         "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",

diff --git a/pkg/kv/kvserver/BUILD.bazel b/pkg/kv/kvserver/BUILD.bazel
@@ -463,6 +463,7 @@ go_test(
         "//pkg/kv/kvserver/tscache",
         "//pkg/kv/kvserver/txnwait",
         "//pkg/kv/kvserver/uncertainty",
+        "//pkg/kv/kvtestutils",
         "//pkg/multitenant",
         "//pkg/multitenant/tenantcapabilities/tenantcapabilitiesauthorizer",
         "//pkg/raft",

diff --git a/pkg/kv/kvserver/client_merge_test.go b/pkg/kv/kvserver/client_merge_test.go
@@ -40,6 +40,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/stateloader"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/txnwait"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
@@ -5257,7 +5258,7 @@ func setupClusterWithSubsumedRange(
 		testutils.SucceedsSoon(t, func() error {
 			var err error
 			newDesc, err = tc.AddVoters(desc.StartKey.AsRawKey(), tc.Target(1))
-			if kv.IsExpectedRelocateError(err) {
+			if kvtestutils.IsExpectedRelocateError(err) {
 				// Retry.
 				return errors.Wrap(err, "ChangeReplicas received error")
 			}

diff --git a/pkg/kv/kvserver/client_relocate_range_test.go b/pkg/kv/kvserver/client_relocate_range_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -84,7 +85,7 @@ func requireRelocationFailure(
 			nonVoterTargets,
 			true, /* transferLeaseToFirstVoter */
 		)
-		if kv.IsExpectedRelocateError(err) {
+		if kvtestutils.IsExpectedRelocateError(err) {
 			return err
 		}
 		require.Regexp(t, errRegExp, err)

diff --git a/pkg/kv/kvtestutils/BUILD.bazel b/pkg/kv/kvtestutils/BUILD.bazel
@@ -0,0 +1,14 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "kvtestutils",
+    srcs = ["test_utils.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvtestutils",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/kv/kvbase",
+        "//pkg/testutils",
+        "//pkg/util/tracing",
+        "//pkg/util/tracing/tracingpb",
+    ],
+)
diff --git a/pkg/kv/test_utils.go → pkg/kv/kvtestutils/test_utils.go b/pkg/kv/test_utils.go → pkg/kv/kvtestutils/test_utils.go
@@ -3,7 +3,7 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
-package kv
+package kvtestutils
 
 import (
 	"strings"

diff --git a/pkg/kv/txn_external_test.go b/pkg/kv/txn_external_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -504,7 +505,7 @@ func testTxnNegotiateAndSendDoesNotBlock(t *testing.T, multiRange, strict, route
 					// assertion.
 					rec := collectAndFinish()
 					expFollowerRead := store.StoreID() != lh.StoreID && strict && routeNearest
-					wasFollowerRead := kv.OnlyFollowerReads(rec)
+					wasFollowerRead := kvtestutils.OnlyFollowerReads(rec)
 					ambiguous := !strict && routeNearest
 					if expFollowerRead != wasFollowerRead && !ambiguous {
 						if expFollowerRead {