From 2379d76a7d1ba935ab8571002450f1060a3b65fb Mon Sep 17 00:00:00 2001
From: David Reiss <david@temporal.io>
Date: Fri, 25 Oct 2024 11:27:29 -0700
Subject: [PATCH] shard versioning at test level

---
 tests/namespace_delete_test.go         |  2 +-
 tests/testcore/functional_test_base.go | 25 +++++++++++++++++--------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/tests/namespace_delete_test.go b/tests/namespace_delete_test.go
index 46211ab5229..04ecf459ed9 100644
--- a/tests/namespace_delete_test.go
+++ b/tests/namespace_delete_test.go
@@ -80,7 +80,7 @@ func TestNamespaceSuite(t *testing.T) {
 
 func (s *namespaceTestSuite) SetupSuite() {
 	// This suite doesn't embed FunctionalTestBase so we have to call CheckTestShard manually.
-	testcore.CheckTestShard(s.T())
+	testcore.CheckTestShard(s.T(), false)
 
 	s.logger = log.NewTestLogger()
 	s.testClusterFactory = testcore.NewTestClusterFactory()
diff --git a/tests/testcore/functional_test_base.go b/tests/testcore/functional_test_base.go
index c8fc61c405c..3542536525f 100644
--- a/tests/testcore/functional_test_base.go
+++ b/tests/testcore/functional_test_base.go
@@ -31,7 +31,6 @@ import (
 	"fmt"
 	"maps"
 	"os"
-	"slices"
 	"strconv"
 	"strings"
 	"testing"
@@ -90,9 +89,12 @@ type (
 )
 
 var (
-	smokeTestSuites = []string{
-		"TestAdvancedVisibilitySuite",
-		"TestClientMiscTestSuite",
+	smokeTestSuites = map[string]bool{
+		"TestAdvancedVisibilitySuite": true,
+		"TestClientMiscTestSuite":     true,
+	}
+	shardByTestSuites = map[string]bool{
+		"TestVersioningFunctionalSuite": true,
 	}
 )
 
@@ -164,7 +166,7 @@ func (s *FunctionalTestBase) SetDynamicConfigOverrides(dynamicConfig map[dynamic
 // Furthermore, all test suites in the "tests/" directory that don't inherit
 // from FunctionalTestBase must implement SetupSuite that calls checkTestShard.
 func (s *FunctionalTestBase) SetupSuite(defaultClusterConfigFile string, options ...Option) {
-	CheckTestShard(s.T())
+	CheckTestShard(s.T(), false)
 
 	s.testClusterFactory = NewTestClusterFactory()
 
@@ -215,15 +217,22 @@ func (s *FunctionalTestBase) SetupSuite(defaultClusterConfigFile string, options
 }
 
 func (s *FunctionalTestBase) SetupTest() {
+	CheckTestShard(s.T(), true)
 }
 
 // CheckTestShard supports test sharding based on environment variables.
-func CheckTestShard(t *testing.T) {
+func CheckTestShard(t *testing.T, atTestLevel bool) {
+	suiteName, _, _ := strings.Cut(t.Name(), "/")
+
+	shouldShardAtTestLevel := shardByTestSuites[suiteName]
+	if shouldShardAtTestLevel != atTestLevel {
+		return
+	}
+
 	indexStr := os.Getenv("TEST_SHARD_INDEX")
 	// special value to run only a few suites on the extended set of persistence drivers
 	if indexStr == "smoke_only" {
-		suiteName, _, _ := strings.Cut(t.Name(), "/")
-		if slices.Contains(smokeTestSuites, suiteName) {
+		if smokeTestSuites[suiteName] {
 			t.Logf("Running %s in smoke tests", t.Name())
 			return
 		}