Update shard persistence test (#2928)

* Add new shard persistence test
* Deprecate old shard persistence test
* Restructure existing test

common/persistence/persistence-tests/cassandra_test.go

@@ -44,13 +44,6 @@ func TestCassandraMetadataPersistenceV2(t *testing.T) {
 	suite.Run(t, s)
 }
 
-func TestCassandraShardPersistence(t *testing.T) {
-	s := new(ShardPersistenceSuite)
-	s.TestBase = NewTestBaseWithCassandra(&TestBaseOptions{})
-	s.TestBase.Setup(nil)
-	suite.Run(t, s)
-}
-
 func TestQueuePersistence(t *testing.T) {
 	s := new(QueuePersistenceSuite)
 	s.TestBase = NewTestBaseWithCassandra(&TestBaseOptions{})

common/persistence/persistence-tests/mysql_test.go

@@ -44,13 +44,6 @@ func TestMySQLMetadataPersistenceSuiteV2(t *testing.T) {
 	suite.Run(t, s)
 }
 
-func TestMySQLShardPersistenceSuite(t *testing.T) {
-	s := new(ShardPersistenceSuite)
-	s.TestBase = NewTestBaseWithSQL(GetMySQLTestClusterOption())
-	s.TestBase.Setup(nil)
-	suite.Run(t, s)
-}
-
 func TestMySQLQueuePersistence(t *testing.T) {
 	s := new(QueuePersistenceSuite)
 	s.TestBase = NewTestBaseWithSQL(GetMySQLTestClusterOption())

common/persistence/persistence-tests/postgres_test.go

@@ -44,13 +44,6 @@ func TestPostgreSQLMetadataPersistenceSuiteV2(t *testing.T) {
 	suite.Run(t, s)
 }
 
-func TestPostgreSQLShardPersistenceSuite(t *testing.T) {
-	s := new(ShardPersistenceSuite)
-	s.TestBase = NewTestBaseWithSQL(GetPostgreSQLTestClusterOption())
-	s.TestBase.Setup(nil)
-	suite.Run(t, s)
-}
-
 func TestPostgreSQLClusterMetadataPersistence(t *testing.T) {
 	s := new(ClusterMetadataManagerSuite)
 	s.TestBase = NewTestBaseWithSQL(GetPostgreSQLTestClusterOption())

common/persistence/persistence-tests/sqlite_test.go

@@ -44,13 +44,6 @@ func TestSQLiteMetadataPersistenceSuiteV2(t *testing.T) {
 	suite.Run(t, s)
 }
 
-func TestSQLiteShardPersistenceSuite(t *testing.T) {
-	s := new(ShardPersistenceSuite)
-	s.TestBase = NewTestBaseWithSQL(GetSQLiteMemoryTestClusterOption())
-	s.TestBase.Setup(nil)
-	suite.Run(t, s)
-}
-
 func TestSQLiteClusterMetadataPersistence(t *testing.T) {
 	s := new(ClusterMetadataManagerSuite)
 	s.TestBase = NewTestBaseWithSQL(GetSQLiteMemoryTestClusterOption())
@@ -79,13 +72,6 @@ func TestSQLiteFileMetadataPersistenceSuiteV2(t *testing.T) {
 	suite.Run(t, s)
 }
 
-func TestSQLiteFileShardPersistenceSuite(t *testing.T) {
-	s := new(ShardPersistenceSuite)
-	s.TestBase = NewTestBaseWithSQL(GetSQLiteFileTestClusterOption())
-	s.TestBase.Setup(nil)
-	suite.Run(t, s)
-}
-
 func TestSQLiteFileClusterMetadataPersistence(t *testing.T) {
 	s := new(ClusterMetadataManagerSuite)
 	s.TestBase = NewTestBaseWithSQL(GetSQLiteFileTestClusterOption())

common/persistence/tests/cassandra_test.go

@@ -27,37 +27,28 @@ package tests
 import (
 	"testing"
 
-	"go.uber.org/zap/zaptest"
-
 	"github.com/stretchr/testify/suite"
 
-	"go.temporal.io/server/common/log"
-	"go.temporal.io/server/common/persistence/cassandra"
 	"go.temporal.io/server/common/persistence/serialization"
 	_ "go.temporal.io/server/common/persistence/sql/sqlplugin/mysql"
-	"go.temporal.io/server/common/resolver"
 )
 
-func setUpCassandraTest(t *testing.T) (CassandraTestData, func()) {
-	var testData CassandraTestData
-	testData.Cfg = NewCassandraConfig()
-	testData.Logger = log.NewZapLogger(zaptest.NewLogger(t))
-	SetUpCassandraDatabase(testData.Cfg, testData.Logger)
-	SetUpCassandraSchema(testData.Cfg, testData.Logger)
-
-	testData.Factory = cassandra.NewFactory(
-		*testData.Cfg,
-		resolver.NewNoopResolver(),
-		testCassandraClusterName,
-		testData.Logger,
-	)
+func TestCassandraShardStoreSuite(t *testing.T) {
+	testData, tearDown := setUpCassandraTest(t)
+	defer tearDown()
 
-	tearDown := func() {
-		testData.Factory.Close()
-		TearDownCassandraKeyspace(testData.Cfg)
+	shardStore, err := testData.Factory.NewShardStore()
+	if err != nil {
+		t.Fatalf("unable to create Cassandra DB: %v", err)
 	}
 
-	return testData, tearDown
+	s := NewShardSuite(
+		t,
+		shardStore,
+		serialization.NewSerializer(),
+		testData.Logger,
+	)
+	suite.Run(t, s)
 }
 
 func TestCassandraExecutionMutableStateStoreSuite(t *testing.T) {
@@ -78,7 +69,8 @@ func TestCassandraExecutionMutableStateStoreSuite(t *testing.T) {
 		shardStore,
 		executionStore,
 		serialization.NewSerializer(),
-		testData.Logger)
+		testData.Logger,
+	)
 	suite.Run(t, s)
 }
 

common/persistence/tests/cassandra_test_util.go

@@ -31,8 +31,10 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+	"testing"
 
 	"github.com/blang/semver/v4"
+	"go.uber.org/zap/zaptest"
 
 	"go.temporal.io/server/common/config"
 	"go.temporal.io/server/common/log"
@@ -61,10 +63,34 @@ const (
 	testCassandraDatabaseNameSuffix = "temporal_persistence"
 )
 
-type CassandraTestData struct {
-	Cfg     *config.Cassandra
-	Factory *cassandra.Factory
-	Logger  log.Logger
+type (
+	CassandraTestData struct {
+		Cfg     *config.Cassandra
+		Factory *cassandra.Factory
+		Logger  log.Logger
+	}
+)
+
+func setUpCassandraTest(t *testing.T) (CassandraTestData, func()) {
+	var testData CassandraTestData
+	testData.Cfg = NewCassandraConfig()
+	testData.Logger = log.NewZapLogger(zaptest.NewLogger(t))
+	SetUpCassandraDatabase(testData.Cfg, testData.Logger)
+	SetUpCassandraSchema(testData.Cfg, testData.Logger)
+
+	testData.Factory = cassandra.NewFactory(
+		*testData.Cfg,
+		resolver.NewNoopResolver(),
+		testCassandraClusterName,
+		testData.Logger,
+	)
+
+	tearDown := func() {
+		testData.Factory.Close()
+		TearDownCassandraKeyspace(testData.Cfg)
+	}
+
+	return testData, tearDown
 }
 
 func SetUpCassandraDatabase(cfg *config.Cassandra, logger log.Logger) {