admin_db_clean_cmd test Part II

tools/cli/admin_commands.go

@@ -308,7 +308,7 @@
 	if err != nil {
 		return commoncli.Problem("Error in Admin delete WF: ", err)
 	}
-	exeStore, err := getDeps(c).initializeExecutionStore(c, shardIDInt)
+	exeStore, err := getDeps(c).initializeExecutionManager(c, shardIDInt)
 	if err != nil {
 		return commoncli.Problem("Error in Admin delete WF: ", err)
 	}

tools/cli/admin_commands_test.go

@@ -51,6 +51,7 @@ type cliTestData struct {
 	mockAdminClient    *admin.MockClient
 	ioHandler          *testIOHandler
 	app                *cli.App
+	mockManagerFactory *MockManagerFactory
 }
 
 func newCLITestData(t *testing.T) *cliTestData {
@@ -60,15 +61,17 @@ func newCLITestData(t *testing.T) *cliTestData {
 
 	td.mockFrontendClient = frontend.NewMockClient(ctrl)
 	td.mockAdminClient = admin.NewMockClient(ctrl)
+	td.mockManagerFactory = NewMockManagerFactory(ctrl)
 	td.ioHandler = &testIOHandler{}
 
-	// make a new CLI App with client factory returning our frontend and admin clients
+	// Create a new CLI app with client factory and persistence manager factory
 	td.app = NewCliApp(
 		&clientFactoryMock{
 			serverFrontendClient: td.mockFrontendClient,
 			serverAdminClient:    td.mockAdminClient,
 		},
 		WithIOHandler(td.ioHandler),
+		WithManagerFactory(td.mockManagerFactory), // Inject the mocked persistence manager factory
 	)
 	return &td
 }

tools/cli/admin_db_clean_command.go

@@ -121,7 +121,8 @@
 			continue
 		}
 
-		fmt.Println(string(data))
+		output := getDeps(c).Output()
+		output.Write([]byte(string(data) + "\n"))
 	}
 	return nil
 }
@@ -131,15 +132,19 @@
 	invariants []executions.InvariantFactory,
 	execution *store.ScanOutputEntity,
 ) (invariant.ManagerFixResult, error) {
-	execManager, err := getDeps(c).initializeExecutionStore(c, execution.Execution.(entity.Entity).GetShardID())
-	defer execManager.Close()
+	execManager, err := getDeps(c).initializeExecutionManager(c, execution.Execution.(entity.Entity).GetShardID())
+	if execManager != nil {
+		defer execManager.Close()
+	}
 	if err != nil {
-		return invariant.ManagerFixResult{}, fmt.Errorf("Error in fix execution: %w", err)
+		return invariant.ManagerFixResult{}, err
 	}
 	historyV2Mgr, err := getDeps(c).initializeHistoryManager(c)
-	defer historyV2Mgr.Close()
+	if historyV2Mgr != nil {
+		defer historyV2Mgr.Close()
+	}
 	if err != nil {
-		return invariant.ManagerFixResult{}, fmt.Errorf("Error in fix execution: %w", err)
+		return invariant.ManagerFixResult{}, err
 	}
 	pr := persistence.NewPersistenceRetryer(
 		execManager,
@@ -156,7 +161,12 @@
 	ctx, cancel, err := newContext(c)
 	defer cancel()
 	if err != nil {
-		return invariant.ManagerFixResult{}, fmt.Errorf("Context not created: %w", err)
+		return invariant.ManagerFixResult{}, err
 	}
-	return invariant.NewInvariantManager(ivs).RunFixes(ctx, execution.Execution), nil
+	invariantManager, err := getDeps(c).initializeInvariantManager(ivs)
+	if err != nil {
+		return invariant.ManagerFixResult{}, err
+	}
+
+	return invariantManager.RunFixes(ctx, execution.Execution.(entity.Entity)), nil
 }

tools/cli/admin_db_clean_command_test.go

@@ -24,21 +24,46 @@ package cli
 
 import (
 	"flag"
+	"fmt"
 	"os"
 	"testing"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/urfave/cli/v2"
+
+	"github.com/uber/cadence/common/persistence"
+	"github.com/uber/cadence/common/reconciliation/invariant"
 )
 
-func TestAdminDBClean_errorCases(t *testing.T) {
+func TestAdminDBClean_noFixExecution(t *testing.T) {
 	tests := []struct {
 		name           string
 		setupContext   func(app *cli.App) *cli.Context
 		inputFileData  string // Simulate the content of the input file
 		expectedOutput string
 		expectedError  string
 	}{
+		{
+			name: "SuccessCase_emptyResult",
+			setupContext: func(app *cli.App) *cli.Context {
+				set := flag.NewFlagSet("test", 0)
+				// Define flags
+				set.String(FlagScanType, "", "scan type flag")
+				set.String(FlagInputFile, "", "Input file flag")
+				// Use a StringSliceFlag to simulate multiple collection values
+				set.Var(cli.NewStringSlice("CollectionHistory", "CollectionDomain"), FlagInvariantCollection, "invariant collection flag")
+
+				// Set actual values for the flags
+				_ = set.Set(FlagScanType, "ConcreteExecutionType")
+				_ = set.Set(FlagInputFile, "input.json")
+
+				return cli.NewContext(app, set, nil)
+			},
+			inputFileData:  `[{"Execution": {"ShardID": 1}}]`, // Simulate the content of input file
+			expectedOutput: ``,
+			expectedError:  "",
+		},
 		{
 			name: "MissingRequiredFlagScanType",
 			setupContext: func(app *cli.App) *cli.Context {
@@ -144,3 +169,152 @@ func TestAdminDBClean_errorCases(t *testing.T) {
 		})
 	}
 }
+
+func TestAdminDBClean_inFixExecution(t *testing.T) {
+	tests := []struct {
+		name           string
+		contextSetup   func(td *cliTestData, inputFilePath string) *cli.Context
+		mockSetup      func(td *cliTestData)
+		inputFileData  string
+		expectedError  string
+		expectedOutput string
+	}{
+		{
+			name: "Success",
+			contextSetup: func(td *cliTestData, inputFilePath string) *cli.Context {
+				set := flag.NewFlagSet("test", 0)
+				set.String(FlagScanType, "", "scan type flag")
+				set.String(FlagInputFile, "", "Input file flag")
+				set.Var(cli.NewStringSlice("CollectionHistory", "CollectionDomain"), FlagInvariantCollection, "invariant collection flag")
+
+				_ = set.Set(FlagScanType, "ConcreteExecutionType")
+				_ = set.Set(FlagInputFile, inputFilePath) // 传递临时文件路径
+
+				return cli.NewContext(td.app, set, nil)
+			},
+			mockSetup: func(td *cliTestData) {
+				ctrl := gomock.NewController(t)
+				mockExecManager := persistence.NewMockExecutionManager(ctrl)
+				mockHistoryManager := persistence.NewMockHistoryManager(ctrl)
+				mockInvariantManager := invariant.NewMockManager(ctrl)
+
+				mockExecManager.EXPECT().Close().Times(1)
+				mockHistoryManager.EXPECT().Close().Times(1)
+
+				mockInvariantManager.EXPECT().RunFixes(gomock.Any(), gomock.Any()).Return(invariant.ManagerFixResult{}).Times(1)
+
+				td.mockManagerFactory.EXPECT().initializeExecutionManager(gomock.Any(), gomock.Any()).Return(mockExecManager, nil).Times(1)
+				td.mockManagerFactory.EXPECT().initializeHistoryManager(gomock.Any()).Return(mockHistoryManager, nil).Times(1)
+				td.mockManagerFactory.EXPECT().initializeInvariantManager(gomock.Any()).Return(mockInvariantManager, nil).Times(1)
+
+			},
+			inputFileData: `{"Execution": {"ShardID": 1}, "Result": {}}`,
+			expectedOutput: `{"Execution":{"BranchToken":null,"TreeID":"","BranchID":"","ShardID":1,"DomainID":"","WorkflowID":"","RunID":"","State":0},"Input":{"Execution":{"BranchToken":null,"TreeID":"","BranchID":"","ShardID":1,"DomainID":"","WorkflowID":"","RunID":"","State":0},"Result":{"CheckResultType":"","DeterminingInvariantType":null,"CheckResults":null}},"Result":{"FixResultType":"","DeterminingInvariantName":null,"FixResults":null}}
+`,
+			expectedError: "",
+		},
+		{
+			name: "init invariant manager error",
+			contextSetup: func(td *cliTestData, inputFilePath string) *cli.Context {
+				set := flag.NewFlagSet("test", 0)
+				set.String(FlagScanType, "", "scan type flag")
+				set.String(FlagInputFile, "", "Input file flag")
+				set.Var(cli.NewStringSlice("CollectionHistory", "CollectionDomain"), FlagInvariantCollection, "invariant collection flag")
+
+				_ = set.Set(FlagScanType, "ConcreteExecutionType")
+				_ = set.Set(FlagInputFile, inputFilePath) // 传递临时文件路径
+
+				return cli.NewContext(td.app, set, nil)
+			},
+			mockSetup: func(td *cliTestData) {
+				ctrl := gomock.NewController(t)
+				mockExecManager := persistence.NewMockExecutionManager(ctrl)
+				mockHistoryManager := persistence.NewMockHistoryManager(ctrl)
+
+				mockExecManager.EXPECT().Close().Times(1)
+				mockHistoryManager.EXPECT().Close().Times(1)
+
+				td.mockManagerFactory.EXPECT().initializeExecutionManager(gomock.Any(), gomock.Any()).Return(mockExecManager, nil).Times(1)
+				td.mockManagerFactory.EXPECT().initializeHistoryManager(gomock.Any()).Return(mockHistoryManager, nil).Times(1)
+				td.mockManagerFactory.EXPECT().initializeInvariantManager(gomock.Any()).Return(nil, fmt.Errorf("init invariant manager error")).Times(1)
+
+			},
+			inputFileData:  `{"Execution": {"ShardID": 1}, "Result": {}}`,
+			expectedOutput: ``,
+			expectedError:  "Error in fix execution: : init invariant manager error",
+		},
+		{
+			name: "init execution manager error",
+			contextSetup: func(td *cliTestData, inputFilePath string) *cli.Context {
+				set := flag.NewFlagSet("test", 0)
+				set.String(FlagScanType, "", "scan type flag")
+				set.String(FlagInputFile, "", "Input file flag")
+				set.Var(cli.NewStringSlice("CollectionHistory", "CollectionDomain"), FlagInvariantCollection, "invariant collection flag")
+
+				_ = set.Set(FlagScanType, "ConcreteExecutionType")
+				_ = set.Set(FlagInputFile, inputFilePath) // 传递临时文件路径
+
+				return cli.NewContext(td.app, set, nil)
+			},
+			mockSetup: func(td *cliTestData) {
+				td.mockManagerFactory.EXPECT().initializeExecutionManager(gomock.Any(), gomock.Any()).Return(nil, fmt.Errorf("init execution manager error")).Times(1)
+			},
+			inputFileData:  `{"Execution": {"ShardID": 1}, "Result": {}}`,
+			expectedOutput: ``,
+			expectedError:  "Error in fix execution: : init execution manager error",
+		},
+		{
+			name: "init history manager error",
+			contextSetup: func(td *cliTestData, inputFilePath string) *cli.Context {
+				set := flag.NewFlagSet("test", 0)
+				set.String(FlagScanType, "", "scan type flag")
+				set.String(FlagInputFile, "", "Input file flag")
+				set.Var(cli.NewStringSlice("CollectionHistory", "CollectionDomain"), FlagInvariantCollection, "invariant collection flag")
+
+				_ = set.Set(FlagScanType, "ConcreteExecutionType")
+				_ = set.Set(FlagInputFile, inputFilePath) // 传递临时文件路径
+
+				return cli.NewContext(td.app, set, nil)
+			},
+			mockSetup: func(td *cliTestData) {
+				ctrl := gomock.NewController(t)
+				mockExecManager := persistence.NewMockExecutionManager(ctrl)
+				mockExecManager.EXPECT().Close().Times(1)
+
+				td.mockManagerFactory.EXPECT().initializeExecutionManager(gomock.Any(), gomock.Any()).Return(mockExecManager, nil).Times(1)
+				td.mockManagerFactory.EXPECT().initializeHistoryManager(gomock.Any()).Return(nil, fmt.Errorf("init history manager error")).Times(1)
+			},
+			inputFileData:  `{"Execution": {"ShardID": 1}, "Result": {}}`,
+			expectedOutput: ``,
+			expectedError:  "Error in fix execution: : init history manager error",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			td := newCLITestData(t)
+			tt.mockSetup(td)
+
+			inputFile, err := os.CreateTemp("", "test_input_*.json")
+			assert.NoError(t, err)
+			defer os.Remove(inputFile.Name())
+
+			if tt.inputFileData != "" {
+				_, err = inputFile.WriteString(tt.inputFileData)
+				assert.NoError(t, err)
+			}
+			inputFile.Close()
+
+			c := tt.contextSetup(td, inputFile.Name())
+
+			err = AdminDBClean(c)
+			if tt.expectedError != "" {
+				assert.ErrorContains(t, err, tt.expectedError)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			assert.Equal(t, tt.expectedOutput, td.consoleOutput())
+		})
+	}
+}

tools/cli/admin_db_scan_command.go

@@ -139,7 +139,7 @@
 	invariants []executions.InvariantFactory,
 	fetcher executions.ExecutionFetcher,
 ) (interface{}, invariant.ManagerCheckResult, error) {
-	execManager, err := getDeps(c).initializeExecutionStore(c, common.WorkflowIDToHistoryShard(req.WorkflowID, numberOfShards))
+	execManager, err := getDeps(c).initializeExecutionManager(c, common.WorkflowIDToHistoryShard(req.WorkflowID, numberOfShards))
 	defer execManager.Close()
 	if err != nil {
 		return nil, invariant.ManagerCheckResult{}, fmt.Errorf("Error in execution check: %w", err)
@@ -200,7 +200,7 @@
 	outputFile *os.File,
 ) error {
 
-	client, err := getDeps(c).initializeExecutionStore(c, shardID)
+	client, err := getDeps(c).initializeExecutionManager(c, shardID)
 	defer client.Close()
 	if err != nil {
 		commoncli.Problem("Error in Admin DB unsupported WF scan: ", err)

tools/cli/admin_timers.go

@@ -80,7 +80,7 @@
 		return nil, fmt.Errorf("error in NewDBLoadCloser: failed to get shard ID: %w", err)
 	}
 
-	executionManager, err := getDeps(c).initializeExecutionStore(c, shardID)
+	executionManager, err := getDeps(c).initializeExecutionManager(c, shardID)
 	if err != nil {
 		return nil, fmt.Errorf("error in NewDBLoadCloser: failed to initialize execution store: %w", err)
 	}

tools/cli/app.go

@@ -50,6 +50,22 @@
 	}
 }
 
+// WithManagerFactory sets the ManagerFactory for the CLI app.
+func WithManagerFactory(factory ManagerFactory) CLIAppOptions {
+	return func(app *cli.App) {
+		if app.Metadata == nil {
+			return
+		}
+
+		d, ok := app.Metadata[depsKey].(*deps)
+		if !ok {
+			return
+		}
+
+		d.ManagerFactory = factory
+	}
+}
+
 // NewCliApp instantiates a new instance of the CLI application
 func NewCliApp(cf ClientFactory, opts ...CLIAppOptions) *cli.App {
 	version := fmt.Sprintf("CLI feature version: %v \n"+
@@ -63,7 +79,7 @@
 	app.Usage = "A command-line tool for cadence users"
 	app.Version = version
 	app.Metadata = map[string]any{
-		depsKey: &deps{ClientFactory: cf, IOHandler: &defaultIOHandler{app: app}, PersistenceManagerFactory: &defaultPersistenceManagerFactory{}},
+		depsKey: &deps{ClientFactory: cf, IOHandler: &defaultIOHandler{app: app}, ManagerFactory: &defaultManagerFactory{}},
 	}
 	app.Flags = []cli.Flag{
 		&cli.StringFlag{
@@ -255,7 +271,7 @@
 type cliDeps interface {
 	ClientFactory
 	IOHandler
-	PersistenceManagerFactory
+	ManagerFactory
 }
 
 type IOHandler interface {
@@ -306,5 +322,5 @@
 type deps struct {
 	ClientFactory
 	IOHandler
-	PersistenceManagerFactory
+	ManagerFactory
 }