Adding tests for WorkflowQuery as a Nexus Op

internal/workertest/env.go

@@ -106,6 +106,10 @@ func (env *TestEnvironment) NexusEndpointName() string {
 	return env.nexusEndpointName
 }
 
+func (env *TestEnvironment) Namespace() string {
+	return testNamespace
+}
+
 func SetupTestEnvironment(t *testing.T, opts ...TestEnvOption) *TestEnvironment {
 	cfg := testEnvConfig{
 		executorTimeout: defaultTestRunTimeout,

internal/workertest/historyrequire.go

@@ -262,6 +262,8 @@ func mapIsSuperset(actual, expected map[string]any) bool {
 	return true
 }
 
+// looselyEqual is used by [mapIsSuperset] and in turn for
+// partial/full history matching
 func looselyEqual(x, y any) bool {
 	switch x := x.(type) {
 	case float64:
@@ -278,6 +280,20 @@ func looselyEqual(x, y any) bool {
 			return mapIsSuperset(x, yMap)
 		}
 		return false
+	case []any:
+		// Match the unordered expected list so a spec can
+		// assert only what it knows. Eg, when asserting links, if runID is
+		// not known, it should be possible to assert on just WID+NS
+		yList, ok := y.([]any)
+		if !ok || len(yList) != len(x) {
+			return false
+		}
+		for i, yv := range yList {
+			if !looselyEqual(x[i], yv) {
+				return false
+			}
+		}
+		return true
 	}
 	return reflect.DeepEqual(x, y)
 }

loadgen/kitchen_sink_executor_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"github.com/temporalio/omes/clioptions"
 	. "github.com/temporalio/omes/internal/workertest"
@@ -18,13 +19,11 @@ import (
 	"go.temporal.io/api/enums/v1"
 	"go.temporal.io/api/history/v1"
 	"go.temporal.io/api/workflowservice/v1"
-	"go.temporal.io/sdk/client"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/emptypb"
 )
 
-const namespace = "default"
-
 var (
 	sdks = []clioptions.Language{
 		clioptions.LangGo,
@@ -59,6 +58,14 @@ var (
 		clioptions.LangTypeScript: "dostandaloneactivityoperatorcommands is not supported",
 		clioptions.LangDotNet:     "dostandaloneactivityoperatorcommands is not supported",
 	}
+
+	nexusQueryUnsupportedSDKs = map[clioptions.Language]string{
+		clioptions.LangPython:     "query-workflow",
+		clioptions.LangJava:       "executenexusoperation is not supported",
+		clioptions.LangRuby:       "executenexusoperation is not supported",
+		clioptions.LangTypeScript: "executenexusoperation is not supported",
+		clioptions.LangDotNet:     "executenexusoperation is not supported",
+	}
 )
 
 type testCase struct {
@@ -67,6 +74,17 @@ type testCase struct {
 	historyMatcher          HistoryMatcher
 	expectedUnsupportedErrs map[clioptions.Language]string
 	expectedWorkflowError   string
+	populateFn              func(t *testing.T) testCase
+}
+
+// populate creates per-run test data for cases that require it.
+func (tc testCase) populate(t *testing.T) testCase {
+	t.Helper() // Keeps error lines pointing to the actual test table
+
+	if tc.populateFn == nil {
+		return tc
+	}
+	return tc.populateFn(t)
 }
 
 // TestKitchenSink tests specific kitchensink features across SDKs.
@@ -939,6 +957,10 @@ func TestKitchenSink(t *testing.T) {
 				NexusOperationCompleted`),
 			expectedUnsupportedErrs: nexusUnsupportedSDKs,
 		},
+		{
+			name:       "NexusOperation/QueryWorkflow",
+			populateFn: populateQueryWorkflowTestCase,
+		},
 		{
 			name: "NexusOperation/Async",
 			testInput: &TestInput{
@@ -1129,13 +1151,6 @@ func TestKitchenSink(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			// Ensure the workflow completes by appending a return action at the end.
-			input := tc.testInput
-			if input.WorkflowInput == nil {
-				input.WorkflowInput = &WorkflowInput{}
-			}
-			input.WorkflowInput.InitialActions = append(input.WorkflowInput.InitialActions, ListActionSet(NewEmptyReturnResultAction())...)
-
 			for _, sdk := range sdks {
 				if onlySDK != "" && string(sdk) != onlySDK {
 					continue // not using t.Skip as it's too noisy
@@ -1149,6 +1164,62 @@ func TestKitchenSink(t *testing.T) {
 	}
 }
 
+func populateQueryWorkflowTestCase(t *testing.T) testCase {
+	t.Helper()
+
+	handlerWorkflowID := fmt.Sprintf("nexus-query-target-%s", uuid.NewString())
+	return testCase{
+		name: "NexusOperation/QueryWorkflow",
+		testInput: &TestInput{
+			WorkflowInput: &WorkflowInput{
+				InitialActions: ListActionSet(
+					// Start a Nexus operation and wait for its signal as a scaffold to
+					// query a running workflow.
+					&Action{
+						Variant: &Action_NexusOperation{
+							NexusOperation: &ExecuteNexusOperation{
+								Operation:                       "echo-async",
+								HandlerWorkflowId:               handlerWorkflowID,
+								HandlerWorkflowIdConflictPolicy: enums.WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING,
+								BeforeActions: ListActionSet(
+									NewSetWorkflowStateAction("query-result", "query successful"),
+								),
+								WaitForSignal: true,
+								AwaitableChoice: &AwaitableChoice{
+									Condition: &AwaitableChoice_WaitStarted{
+										WaitStarted: &emptypb.Empty{},
+									},
+								},
+							},
+						},
+					},
+					&Action{
+						Variant: &Action_NexusOperation{
+							NexusOperation: &ExecuteNexusOperation{
+								Operation: "query-workflow",
+								Input: QueryWorkflowNexusInput(QueryWorkflowTarget{
+									WorkflowID: handlerWorkflowID,
+								}),
+								ExpectedOutput: "query successful",
+								AwaitableChoice: &AwaitableChoice{
+									Condition: &AwaitableChoice_WaitFinish{
+										WaitFinish: &emptypb.Empty{},
+									},
+								},
+							},
+						},
+					},
+				),
+			},
+		},
+		historyMatcher: PartialHistoryMatcher(fmt.Sprintf(`
+			NexusOperationScheduled {"operation":"query-workflow"}
+			NexusOperationCompleted {"links":[{"workflow":{"workflowId":%q,"reason":"Query processed"}}]}`,
+			handlerWorkflowID)),
+		expectedUnsupportedErrs: nexusQueryUnsupportedSDKs,
+	}
+}
+
 func standaloneActivityOperatorCommandsTestCase(
 	name string,
 	commandType DoStandaloneActivityOperatorCommands_CommandType,
@@ -1198,6 +1269,8 @@ func testForSDK(
 	env *TestEnvironment,
 	workflowTimeout time.Duration,
 ) {
+	tc = tc.populate(t)
+
 	// Use mutex to ensure only one Java test runs at a time/a Gradle limitation.
 	if sdk == clioptions.LangJava {
 		javaMutex.Lock()
@@ -1220,8 +1293,18 @@ func testForSDK(
 	// worker picks them up.
 	runTaskQueue := TaskQueueForRun(scenarioInfo.RunID)
 
+	// PrepareTestInput fills the per-run endpoint and task queue into the input in
+	// place, so work on a copy before making further per-run changes.
+	testInput := proto.Clone(tc.testInput).(*TestInput)
+
+	// Ensure the workflow completes by appending a return action at the end.
+	if testInput.WorkflowInput == nil {
+		testInput.WorkflowInput = &WorkflowInput{}
+	}
+	testInput.WorkflowInput.InitialActions = append(testInput.WorkflowInput.InitialActions, ListActionSet(NewEmptyReturnResultAction())...)
+
 	executor := &KitchenSinkExecutor{
-		TestInput: tc.testInput,
+		TestInput: testInput,
 		PrepareTestInput: func(_ context.Context, _ ScenarioInfo, input *TestInput) error {
 			if input.WorkflowInput != nil {
 				for _, actionSet := range input.WorkflowInput.InitialActions {
@@ -1298,7 +1381,7 @@ func testSupportedFeature(
 	_, execErr := env.RunExecutorTest(t, testExecutor, scenarioInfo, sdk)
 
 	taskQueueName := TaskQueueForRun(scenarioInfo.RunID)
-	historyEvents, historyErr := getWorkflowHistory(t, taskQueueName, env.TemporalClient())
+	historyEvents, historyErr := getWorkflowHistory(t, taskQueueName, env)
 	if execErr != nil {
 		if len(historyEvents) > 0 {
 			t.Logf("History events for debugging:")
@@ -1340,10 +1423,11 @@ func (w *kitchenSinkTestWrapper) Run(ctx context.Context, info ScenarioInfo) err
 	return w.executor.Run(ctx, info)
 }
 
-func getWorkflowHistory(t *testing.T, taskQueueName string, temporalClient client.Client) ([]*history.HistoryEvent, error) {
+func getWorkflowHistory(t *testing.T, taskQueueName string, env *TestEnvironment) ([]*history.HistoryEvent, error) {
+	temporalClient := env.TemporalClient()
 	executions, err := temporalClient.ListWorkflow(t.Context(),
 		&workflowservice.ListWorkflowExecutionsRequest{
-			Namespace: namespace,
+			Namespace: env.Namespace(),
 			Query:     fmt.Sprintf("TaskQueue = '%s' AND WorkflowType = 'kitchenSink'", taskQueueName),
 		})
 	if err != nil {

loadgen/kitchensink/helpers.go

@@ -1,6 +1,7 @@
 package kitchensink
 
 import (
+	"encoding/json"
 	"fmt"
 	"math/rand"
 	"time"
@@ -13,6 +14,21 @@ import (
 	"google.golang.org/protobuf/types/known/emptypb"
 )
 
+// QueryWorkflowTarget identifies the workflow that the "query-workflow" Nexus operation queries.
+// Only WorkflowID is required
+type QueryWorkflowTarget struct {
+	WorkflowID string `json:"workflow_id"`
+}
+
+// QueryWorkflowNexusInput encodes a target as the input of a "query-workflow" Nexus operation.
+func QueryWorkflowNexusInput(target QueryWorkflowTarget) string {
+	encoded, err := json.Marshal(target)
+	if err != nil {
+		panic(err)
+	}
+	return string(encoded)
+}
+
 // Using human-readable JSON encoding for payloads to aid with debugging.
 var jsonPayloadConverter = converter.NewProtoJSONPayloadConverter()
 

mise.toml

@@ -25,7 +25,7 @@ ruby = "1.6.0"
 typescript = "1.22.0"
 
 [_.server]
-ref = "abcb94873061dd55e80745cbcc84037abdf22d53"
+ref = "042505361519608454d4a016b3d0971c81c724c8"
 
 [tasks."sync-sdk:go"]
 description = "Sync Go SDK dependencies to the version in mise.toml"

workers/go/apps/lambda/worker.go

@@ -116,7 +116,12 @@ func configureLambdaWorker(opts *lambdaworker.Options) error {
 	ebbFlowActivities := ebbandflow.Activities{}
 
 	service := nexus.NewService(kitchensink.KitchenSinkServiceName)
-	for _, op := range []nexus.RegisterableOperation{kitchensink.EchoSyncOperation, kitchensink.EchoAsyncOperation, kitchensink.StandaloneActivityNexusOperation} {
+	for _, op := range []nexus.RegisterableOperation{
+		kitchensink.EchoSyncOperation,
+		kitchensink.EchoAsyncOperation,
+		kitchensink.StandaloneActivityNexusOperation,
+		kitchensink.QueryWorkflowOperation,
+	} {
 		if err := service.Register(op); err != nil {
 			return fmt.Errorf("failed to register nexus operation: %w", err)
 		}

workers/go/apps/worker/worker.go

@@ -35,6 +35,7 @@ func buildWorker(client sdkclient.Client, context harness.WorkerContext) sdkwork
 		kitchensink.EchoSyncOperation,
 		kitchensink.EchoAsyncOperation,
 		kitchensink.StandaloneActivityNexusOperation,
+		kitchensink.QueryWorkflowOperation,
 	} {
 		if err := service.Register(op); err != nil {
 			panic(err)