diff --git a/common/metrics/metricstest/capture_handler.go b/common/metrics/metricstest/capture_handler.go index 862fb02998..c3b6382969 100644 --- a/common/metrics/metricstest/capture_handler.go +++ b/common/metrics/metricstest/capture_handler.go @@ -1,6 +1,8 @@ package metricstest import ( + "maps" + "slices" "sync" "sync/atomic" "time" @@ -18,19 +20,19 @@ type CapturedRecording struct { // Capture is a specific capture instance. type Capture struct { - recordings map[string][]*CapturedRecording + recordings CaptureSnapshot recordingsLock sync.RWMutex } +type CaptureSnapshot = map[string][]*CapturedRecording + // Snapshot returns a copy of all metrics recorded, keyed by name. -func (c *Capture) Snapshot() map[string][]*CapturedRecording { +func (c *Capture) Snapshot() CaptureSnapshot { c.recordingsLock.RLock() defer c.recordingsLock.RUnlock() - ret := make(map[string][]*CapturedRecording, len(c.recordings)) - for k, v := range c.recordings { - recs := make([]*CapturedRecording, len(v)) - copy(recs, v) - ret[k] = recs + ret := maps.Clone(c.recordings) + for k, v := range ret { + ret[k] = slices.Clone(v) } return ret } @@ -63,7 +65,7 @@ func NewCaptureHandler() *CaptureHandler { // StartCapture returns a started capture. StopCapture should be called on // complete. func (c *CaptureHandler) StartCapture() *Capture { - capture := &Capture{recordings: map[string][]*CapturedRecording{}} + capture := &Capture{recordings: make(CaptureSnapshot)} c.capturesLock.Lock() defer c.capturesLock.Unlock() diff --git a/common/testing/taskpoller/taskpoller.go b/common/testing/taskpoller/taskpoller.go index fd975933fd..7bb9c72cfc 100644 --- a/common/testing/taskpoller/taskpoller.go +++ b/common/testing/taskpoller/taskpoller.go @@ -9,6 +9,7 @@ import ( "github.com/nexus-rpc/sdk-go/nexus" nexuspb "go.temporal.io/api/nexus/v1" + commandpb "go.temporal.io/api/command/v1" enumspb "go.temporal.io/api/enums/v1" failurepb "go.temporal.io/api/failure/v1" historypb "go.temporal.io/api/history/v1" @@ -55,6 +56,17 @@ var ( DrainWorkflowTask = func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return &workflowservice.RespondWorkflowTaskCompletedRequest{}, nil } + // CompleteWorkflowHandler is a workflow task handler that always completes the workflow. + CompleteWorkflowHandler = func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { + return &workflowservice.RespondWorkflowTaskCompletedRequest{ + Commands: []*commandpb.Command{{ + CommandType: enumspb.COMMAND_TYPE_COMPLETE_WORKFLOW_EXECUTION, + Attributes: &commandpb.Command_CompleteWorkflowExecutionCommandAttributes{ + CompleteWorkflowExecutionCommandAttributes: &commandpb.CompleteWorkflowExecutionCommandAttributes{}, + }, + }}, + }, nil + } // CompleteActivityTask returns a RespondActivityTaskCompletedRequest with an auto-generated `Result` from `tv.Any().Payloads()`. CompleteActivityTask = func(tv *testvars.TestVars) func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { return func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { diff --git a/tests/testcore/test_env.go b/tests/testcore/test_env.go index 253dbe1ec1..9621277198 100644 --- a/tests/testcore/test_env.go +++ b/tests/testcore/test_env.go @@ -17,6 +17,7 @@ import ( "go.temporal.io/api/workflowservice/v1" sdkclient "go.temporal.io/sdk/client" sdkworker "go.temporal.io/sdk/worker" + "go.temporal.io/server/api/adminservice/v1" "go.temporal.io/server/common/debug" "go.temporal.io/server/common/dynamicconfig" "go.temporal.io/server/common/log" @@ -45,6 +46,7 @@ type Env interface { Namespace() namespace.Name NamespaceID() namespace.ID FrontendClient() workflowservice.WorkflowServiceClient + AdminClient() adminservice.AdminServiceClient GetTestCluster() *TestCluster CloseShard(namespaceID string, workflowID string) OverrideDynamicConfig(setting dynamicconfig.GenericSetting, value any) (cleanup func())