Update nexus_test_base.go to use a testcore.TestEnv base

Author: long-nt-tran

tests/nexus_test_base.go

@@ -3,9 +3,11 @@ package tests
 import (
 	"context"
 	"errors"
+	"testing"
 
 	"github.com/google/uuid"
 	"github.com/nexus-rpc/sdk-go/nexus"
+	"github.com/stretchr/testify/require"
 	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	nexuspb "go.temporal.io/api/nexus/v1"
@@ -15,29 +17,44 @@ import (
 	"go.temporal.io/api/workflowservice/v1"
 	cnexus "go.temporal.io/server/common/nexus"
 	"go.temporal.io/server/common/nexus/nexusrpc"
+	"go.temporal.io/server/components/nexusoperations"
 	"go.temporal.io/server/tests/testcore"
 )
 
-type NexusTestBaseSuite struct {
-	testcore.FunctionalTestBase
+type NexusTestEnv struct {
+	*testcore.TestEnv
 	useTemporalFailures bool
 }
 
-func (s *NexusTestBaseSuite) createNexusEndpoint(name string, taskQueue string) *nexuspb.Endpoint {
-	resp, err := s.OperatorClient().CreateNexusEndpoint(testcore.NewContext(), &operatorservice.CreateNexusEndpointRequest{
+func newNexusTestEnv(t *testing.T, useTemporalFailures bool, opts ...testcore.TestOption) *NexusTestEnv {
+	env := &NexusTestEnv{
+		TestEnv:             testcore.NewEnv(t, opts...),
+		useTemporalFailures: useTemporalFailures,
+	}
+	if testcore.IsDedicatedCluster(opts) {
+		env.OverrideDynamicConfig(
+			nexusoperations.CallbackURLTemplate,
+			"http://"+env.HttpAPIAddress()+"/nexus/callback",
+		)
+	}
+	return env
+}
+
+func (env *NexusTestEnv) createNexusEndpoint(name string, taskQueue string, t *testing.T) *nexuspb.Endpoint {
+	resp, err := env.OperatorClient().CreateNexusEndpoint(testcore.NewContext(), &operatorservice.CreateNexusEndpointRequest{
 		Spec: &nexuspb.EndpointSpec{
 			Name: name,
 			Target: &nexuspb.EndpointTarget{
 				Variant: &nexuspb.EndpointTarget_Worker_{
 					Worker: &nexuspb.EndpointTarget_Worker{
-						Namespace: s.Namespace().String(),
+						Namespace: env.Namespace().String(),
 						TaskQueue: taskQueue,
 					},
 				},
 			},
 		},
 	})
-	s.NoError(err)
+	require.NoError(t, err)
 	return resp.Endpoint
 }
 
@@ -56,28 +73,28 @@ type nexusTaskResponse struct {
 
 type nexusTaskHandler func(res *workflowservice.PollNexusTaskQueueResponse) (*nexusTaskResponse, error)
 
-func (s *NexusTestBaseSuite) nexusTaskPoller(ctx context.Context, taskQueue string, handler nexusTaskHandler) <-chan error {
-	return s.versionedNexusTaskPoller(ctx, taskQueue, "", handler)
+func (env *NexusTestEnv) nexusTaskPoller(ctx context.Context, taskQueue string, handler nexusTaskHandler) <-chan error {
+	return env.versionedNexusTaskPoller(ctx, taskQueue, "", handler)
 }
 
-func (s *NexusTestBaseSuite) versionedNexusTaskPoller(ctx context.Context, taskQueue, buildID string, handler nexusTaskHandler) <-chan error {
+func (env *NexusTestEnv) versionedNexusTaskPoller(ctx context.Context, taskQueue, buildID string, handler nexusTaskHandler) <-chan error {
 	errCh := make(chan error, 1)
 	go func() {
-		errCh <- s.versionedNexusTaskPollerDo(ctx, taskQueue, buildID, handler)
+		errCh <- env.versionedNexusTaskPollerDo(ctx, taskQueue, buildID, handler)
 	}()
 	return errCh
 }
 
-func (s *NexusTestBaseSuite) versionedNexusTaskPollerDo(ctx context.Context, taskQueue, buildID string, handler nexusTaskHandler) error {
+func (env *NexusTestEnv) versionedNexusTaskPollerDo(ctx context.Context, taskQueue, buildID string, handler nexusTaskHandler) error {
 	var vc *commonpb.WorkerVersionCapabilities
 	if buildID != "" {
 		vc = &commonpb.WorkerVersionCapabilities{
 			BuildId:       buildID,
 			UseVersioning: true,
 		}
 	}
-	res, err := s.GetTestCluster().FrontendClient().PollNexusTaskQueue(ctx, &workflowservice.PollNexusTaskQueueRequest{
-		Namespace: s.Namespace().String(),
+	res, err := env.GetTestCluster().FrontendClient().PollNexusTaskQueue(ctx, &workflowservice.PollNexusTaskQueueRequest{
+		Namespace: env.Namespace().String(),
 		Identity:  uuid.NewString(),
 		TaskQueue: &taskqueuepb.TaskQueue{
 			Name: taskQueue,
@@ -103,9 +120,9 @@ func (s *NexusTestBaseSuite) versionedNexusTaskPollerDo(ctx context.Context, tas
 		var opErr *nexus.OperationError
 		var he *nexus.HandlerError
 		if errors.As(handlerErr, &opErr) {
-			return s.respondNexusTaskCompletedWithOperationError(ctx, res.TaskToken, opErr)
+			return env.respondNexusTaskCompletedWithOperationError(ctx, res.TaskToken, opErr)
 		} else if errors.As(handlerErr, &he) {
-			return s.respondNexusTaskFailed(ctx, res.TaskToken, he)
+			return env.respondNexusTaskFailed(ctx, res.TaskToken, he)
 		}
 		return handlerErr
 	}
@@ -163,8 +180,8 @@ func (s *NexusTestBaseSuite) versionedNexusTaskPollerDo(ctx context.Context, tas
 			panic("unreachable") // nolint:revive // all implementations of HandlerStartOperationResult must be covered here, so this should be unreachable.
 		}
 	}
-	_, err = s.GetTestCluster().FrontendClient().RespondNexusTaskCompleted(ctx, &workflowservice.RespondNexusTaskCompletedRequest{
-		Namespace: s.Namespace().String(),
+	_, err = env.GetTestCluster().FrontendClient().RespondNexusTaskCompleted(ctx, &workflowservice.RespondNexusTaskCompletedRequest{
+		Namespace: env.Namespace().String(),
 		Identity:  uuid.NewString(),
 		TaskToken: res.TaskToken,
 		Response:  response,
@@ -175,8 +192,8 @@ func (s *NexusTestBaseSuite) versionedNexusTaskPollerDo(ctx context.Context, tas
 	return nil
 }
 
-func (s *NexusTestBaseSuite) respondNexusTaskFailed(ctx context.Context, taskToken []byte, he *nexus.HandlerError) error {
-	if s.useTemporalFailures {
+func (env *NexusTestEnv) respondNexusTaskFailed(ctx context.Context, taskToken []byte, he *nexus.HandlerError) error {
+	if env.useTemporalFailures {
 		nexusFailure, err := nexusrpc.DefaultFailureConverter().ErrorToFailure(he)
 		if err != nil {
 			return err
@@ -185,8 +202,8 @@ func (s *NexusTestBaseSuite) respondNexusTaskFailed(ctx context.Context, taskTok
 		if err != nil {
 			return err
 		}
-		_, err = s.GetTestCluster().FrontendClient().RespondNexusTaskFailed(ctx, &workflowservice.RespondNexusTaskFailedRequest{
-			Namespace: s.Namespace().String(),
+		_, err = env.GetTestCluster().FrontendClient().RespondNexusTaskFailed(ctx, &workflowservice.RespondNexusTaskFailedRequest{
+			Namespace: env.Namespace().String(),
 			Identity:  uuid.NewString(),
 			TaskToken: taskToken,
 			Failure:   temporalFailure,
@@ -219,8 +236,8 @@ func (s *NexusTestBaseSuite) respondNexusTaskFailed(ctx context.Context, taskTok
 		protoError.RetryBehavior = enumspb.NEXUS_HANDLER_ERROR_RETRY_BEHAVIOR_NON_RETRYABLE
 	default:
 	}
-	_, err := s.GetTestCluster().FrontendClient().RespondNexusTaskFailed(ctx, &workflowservice.RespondNexusTaskFailedRequest{
-		Namespace: s.Namespace().String(),
+	_, err := env.GetTestCluster().FrontendClient().RespondNexusTaskFailed(ctx, &workflowservice.RespondNexusTaskFailedRequest{
+		Namespace: env.Namespace().String(),
 		Identity:  uuid.NewString(),
 		TaskToken: taskToken,
 		Error:     protoError,
@@ -231,8 +248,8 @@ func (s *NexusTestBaseSuite) respondNexusTaskFailed(ctx context.Context, taskTok
 	return nil
 }
 
-func (s *NexusTestBaseSuite) respondNexusTaskCompletedWithOperationError(ctx context.Context, taskToken []byte, opErr *nexus.OperationError) error {
-	if s.useTemporalFailures {
+func (env *NexusTestEnv) respondNexusTaskCompletedWithOperationError(ctx context.Context, taskToken []byte, opErr *nexus.OperationError) error {
+	if env.useTemporalFailures {
 		nexusFailure, err := nexusrpc.DefaultFailureConverter().ErrorToFailure(opErr)
 		if err != nil {
 			return err
@@ -250,8 +267,8 @@ func (s *NexusTestBaseSuite) respondNexusTaskCompletedWithOperationError(ctx con
 				},
 			},
 		}
-		_, err = s.GetTestCluster().FrontendClient().RespondNexusTaskCompleted(ctx, &workflowservice.RespondNexusTaskCompletedRequest{
-			Namespace: s.Namespace().String(),
+		_, err = env.GetTestCluster().FrontendClient().RespondNexusTaskCompleted(ctx, &workflowservice.RespondNexusTaskCompletedRequest{
+			Namespace: env.Namespace().String(),
 			Identity:  uuid.NewString(),
 			TaskToken: taskToken,
 			Response:  response,
@@ -284,8 +301,8 @@ func (s *NexusTestBaseSuite) respondNexusTaskCompletedWithOperationError(ctx con
 			},
 		},
 	}
-	_, err := s.GetTestCluster().FrontendClient().RespondNexusTaskCompleted(ctx, &workflowservice.RespondNexusTaskCompletedRequest{
-		Namespace: s.Namespace().String(),
+	_, err := env.GetTestCluster().FrontendClient().RespondNexusTaskCompleted(ctx, &workflowservice.RespondNexusTaskCompletedRequest{
+		Namespace: env.Namespace().String(),
 		Identity:  uuid.NewString(),
 		TaskToken: taskToken,
 		Response:  response,

tests/testcore/test_env.go

@@ -358,3 +358,13 @@ func checkTestShard(t *testing.T) {
 	}
 	t.Logf("Running %s in test shard %d/%d", t.Name(), index+1, total)
 }
+
+// IsDedicatedCluster allows caller to inspect whether opts contains
+// testcore.WithDedicatedCluster().
+func IsDedicatedCluster(opts []TestOption) bool {
+	var o testOptions
+	for _, opt := range opts {
+		opt(&o)
+	}
+	return o.dedicatedCluster
+}