From bd9317e2fb7851a85b96c36b9f5f080e027c45a8 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Wed, 12 Jul 2023 16:16:08 -0700
Subject: [PATCH] tenantcostclient: speed up query RU estimate test

This commit speeds up the query RU estimate test in two ways:
- it forces the production values for some of the metamorphic constants
(this brought down the runtime from 93s to 37s for a particular seed)
- it reduces the number of rows inserted from 50k to 20k which should be
(further bringing down the runtime to 13s).

Release note: None
---
 pkg/ccl/multitenantccl/tenantcostclient/BUILD.bazel    |  1 +
 .../tenantcostclient/query_ru_estimate_test.go         | 10 +++++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/pkg/ccl/multitenantccl/tenantcostclient/BUILD.bazel b/pkg/ccl/multitenantccl/tenantcostclient/BUILD.bazel
index b145f0329add..507449e87e93 100644
--- a/pkg/ccl/multitenantccl/tenantcostclient/BUILD.bazel
+++ b/pkg/ccl/multitenantccl/tenantcostclient/BUILD.bazel
@@ -77,6 +77,7 @@ go_test(
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/distsql",
         "//pkg/sql/execinfra",
+        "//pkg/sql/sem/eval",
         "//pkg/sql/sqlliveness",
         "//pkg/sql/sqlliveness/slinstance",
         "//pkg/sql/stats",
diff --git a/pkg/ccl/multitenantccl/tenantcostclient/query_ru_estimate_test.go b/pkg/ccl/multitenantccl/tenantcostclient/query_ru_estimate_test.go
index ec181cd7e1fa..054778d1175f 100644
--- a/pkg/ccl/multitenantccl/tenantcostclient/query_ru_estimate_test.go
+++ b/pkg/ccl/multitenantccl/tenantcostclient/query_ru_estimate_test.go
@@ -23,6 +23,7 @@ import (
 	_ "github.com/cockroachdb/cockroach/pkg/ccl/multitenantccl/tenantcostserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -74,6 +75,13 @@ func TestEstimateQueryRUConsumption(t *testing.T) {
 	tenant1, tenantDB1 := serverutils.StartTenant(t, s, base.TestTenantArgs{
 		TenantID: tenantID,
 		Settings: st,
+		TestingKnobs: base.TestingKnobs{
+			SQLEvalContext: &eval.TestingKnobs{
+				// We disable the randomization of some batch sizes because with
+				// some low values the test takes much longer.
+				ForceProductionValues: true,
+			},
+		},
 	})
 	defer tenant1.Stopper().Stop(ctx)
 	defer tenantDB1.Close()
@@ -87,7 +95,7 @@ func TestEstimateQueryRUConsumption(t *testing.T) {
 	}
 	testCases := []testCase{
 		{ // Insert statement
-			sql:   "INSERT INTO abcd (SELECT t%2, t%3, t, -t FROM generate_series(1,50000) g(t))",
+			sql:   "INSERT INTO abcd (SELECT t%2, t%3, t, -t FROM generate_series(1,20000) g(t))",
 			count: 1,
 		},
 		{ // Point query