mirror of
https://github.com/hanzoai/tasks.git
synced 2026-11-03 03:56:04 +00:00
zeekay
5a601d081b
Full module path rebrand from Temporal to Hanzo Tasks: - go.mod module: github.com/hanzoai/tasks - All internal Go imports updated (1906 files) - cmd/server renamed to cmd/tasksd - Binary name: tasksd (was temporal-server) - Docker images: hanzoai/tasks (was temporalio/server) - README rewritten with Hanzo Tasks branding - LLM.md created with integration notes - Makefile, goreleaser, CI workflows updated - go.temporal.io/sdk and go.temporal.io/api imports preserved (separate repos) Fork of Temporal.io (MIT license).
363 lines
12 KiB
Go
363 lines
12 KiB
Go
package tests
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
"go.temporal.io/api/serviceerror"
|
|
taskqueuepb "go.temporal.io/api/taskqueue/v1"
|
|
workerpb "go.temporal.io/api/worker/v1"
|
|
"go.temporal.io/api/workflowservice/v1"
|
|
"github.com/hanzoai/tasks/common/dynamicconfig"
|
|
"github.com/hanzoai/tasks/common/testing/testvars"
|
|
"github.com/hanzoai/tasks/tests/testcore"
|
|
)
|
|
|
|
type WorkerRegistryTestSuite struct {
|
|
testcore.FunctionalTestBase
|
|
tv *testvars.TestVars
|
|
}
|
|
|
|
func TestWorkerRegistryTestSuite(t *testing.T) {
|
|
s := new(WorkerRegistryTestSuite)
|
|
suite.Run(t, s)
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) SetupTest() {
|
|
s.OverrideDynamicConfig(dynamicconfig.WorkerHeartbeatsEnabled, true)
|
|
s.FunctionalTestBase.SetupTest()
|
|
|
|
s.tv = testvars.New(s.T()).WithTaskQueue(s.TaskQueue()).WithNamespaceName(s.Namespace())
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) TestWorkerRegistry_DescribeWorker() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
// Record heartbeat for 2 workers
|
|
worker1Key := s.tv.WorkerIdentity()
|
|
worker2Key := s.tv.WorkerIdentity() + "_2"
|
|
taskQueue1 := s.tv.WithTaskQueueNumber(1).TaskQueue().Name
|
|
taskQueue2 := s.tv.WithTaskQueueNumber(2).TaskQueue().Name
|
|
|
|
hbResp, err := s.FrontendClient().RecordWorkerHeartbeat(ctx, &workflowservice.RecordWorkerHeartbeatRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerHeartbeat: []*workerpb.WorkerHeartbeat{
|
|
{
|
|
WorkerInstanceKey: worker1Key,
|
|
TaskQueue: taskQueue1,
|
|
TotalStickyCacheHit: 1,
|
|
},
|
|
{
|
|
WorkerInstanceKey: worker2Key,
|
|
TaskQueue: taskQueue2,
|
|
TotalStickyCacheHit: 2,
|
|
},
|
|
},
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(hbResp)
|
|
|
|
// Test error case - worker that doesn't exist
|
|
nonExistentWorkerKey := s.tv.WorkerIdentity() + "_nonexistent"
|
|
_, err = s.FrontendClient().DescribeWorker(ctx, &workflowservice.DescribeWorkerRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerInstanceKey: nonExistentWorkerKey,
|
|
})
|
|
s.Error(err)
|
|
var notFound *serviceerror.NotFound
|
|
s.ErrorAs(err, ¬Found)
|
|
|
|
// Test error case - unknown namespace
|
|
unknownNamespace := s.tv.NamespaceName().String() + "_unknown"
|
|
_, err = s.FrontendClient().DescribeWorker(ctx, &workflowservice.DescribeWorkerRequest{
|
|
Namespace: unknownNamespace,
|
|
WorkerInstanceKey: worker1Key,
|
|
})
|
|
s.Error(err)
|
|
var namespaceNotFound *serviceerror.NamespaceNotFound
|
|
s.ErrorAs(err, &namespaceNotFound)
|
|
|
|
// Test success case - verify worker1 heartbeat data
|
|
{
|
|
resp, err := s.FrontendClient().DescribeWorker(ctx, &workflowservice.DescribeWorkerRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerInstanceKey: worker1Key,
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(resp)
|
|
s.NotNil(resp.GetWorkerInfo())
|
|
|
|
workerHeartbeat := resp.GetWorkerInfo().GetWorkerHeartbeat()
|
|
s.NotNil(workerHeartbeat)
|
|
s.Equal(worker1Key, workerHeartbeat.GetWorkerInstanceKey())
|
|
s.Equal(taskQueue1, workerHeartbeat.GetTaskQueue())
|
|
s.Equal(int32(1), workerHeartbeat.GetTotalStickyCacheHit())
|
|
}
|
|
// Test success case - verify worker2 heartbeat data
|
|
{
|
|
resp, err := s.FrontendClient().DescribeWorker(ctx, &workflowservice.DescribeWorkerRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerInstanceKey: worker2Key,
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(resp)
|
|
s.NotNil(resp.GetWorkerInfo())
|
|
|
|
workerHeartbeat := resp.GetWorkerInfo().GetWorkerHeartbeat()
|
|
s.NotNil(workerHeartbeat)
|
|
s.Equal(worker2Key, workerHeartbeat.GetWorkerInstanceKey())
|
|
s.Equal(taskQueue2, workerHeartbeat.GetTaskQueue())
|
|
s.Equal(int32(2), workerHeartbeat.GetTotalStickyCacheHit())
|
|
}
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) TestWorkerRegistry_ListWorkersWithNoHeartbeats() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: "TaskQueue='no-heartbeats-recorded-on-this-queue'",
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(resp)
|
|
s.Empty(resp.GetWorkersInfo()) //nolint:staticcheck // testing deprecated field
|
|
s.Empty(resp.GetWorkers())
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) TestWorkerRegistry_DescribeWorkerWithNoHeartbeats() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
_, err := s.FrontendClient().DescribeWorker(ctx, &workflowservice.DescribeWorkerRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerInstanceKey: "nonexistent-worker",
|
|
})
|
|
s.Require().Error(err)
|
|
var notFound *serviceerror.NotFound
|
|
var namespaceNotFound *serviceerror.NamespaceNotFound
|
|
s.True(
|
|
errors.As(err, ¬Found) || errors.As(err, &namespaceNotFound),
|
|
"expected NotFound or NamespaceNotFound, got: %v", err,
|
|
)
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) TestWorkerRegistry_ListWorkers() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
// Record heartbeats for 2 workers
|
|
worker1Key := s.tv.WorkerIdentity()
|
|
worker2Key := s.tv.WorkerIdentity() + "_2"
|
|
sharedTaskQueue := s.tv.TaskQueue().Name
|
|
|
|
hbResp, err := s.FrontendClient().RecordWorkerHeartbeat(ctx, &workflowservice.RecordWorkerHeartbeatRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerHeartbeat: []*workerpb.WorkerHeartbeat{
|
|
{
|
|
WorkerInstanceKey: worker1Key,
|
|
TaskQueue: sharedTaskQueue,
|
|
TotalStickyCacheHit: 1,
|
|
},
|
|
{
|
|
WorkerInstanceKey: worker2Key,
|
|
TaskQueue: sharedTaskQueue,
|
|
TotalStickyCacheHit: 2,
|
|
},
|
|
},
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(hbResp)
|
|
|
|
{
|
|
resp, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: fmt.Sprintf("WorkerInstanceKey='%s'", worker1Key),
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(resp)
|
|
s.Len(resp.GetWorkersInfo(), 1)
|
|
|
|
// Verify deprecated WorkersInfo field (backward compatibility)
|
|
workerHeartbeat := resp.GetWorkersInfo()[0].GetWorkerHeartbeat()
|
|
s.Equal(worker1Key, workerHeartbeat.WorkerInstanceKey)
|
|
s.Equal(sharedTaskQueue, workerHeartbeat.TaskQueue)
|
|
s.Equal(int32(1), workerHeartbeat.TotalStickyCacheHit)
|
|
|
|
// Verify new Workers field (WorkerListInfo)
|
|
s.Len(resp.GetWorkers(), 1)
|
|
workerListInfo := resp.GetWorkers()[0]
|
|
s.Equal(worker1Key, workerListInfo.GetWorkerInstanceKey())
|
|
s.Equal(sharedTaskQueue, workerListInfo.GetTaskQueue())
|
|
}
|
|
{
|
|
resp, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: fmt.Sprintf("TaskQueue='%s'", sharedTaskQueue),
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(resp)
|
|
|
|
// Verify deprecated WorkersInfo field
|
|
workers := resp.GetWorkersInfo()
|
|
workersByKey := make(map[string]*workerpb.WorkerHeartbeat)
|
|
for _, workerInfo := range workers {
|
|
heartbeat := workerInfo.GetWorkerHeartbeat()
|
|
workersByKey[heartbeat.WorkerInstanceKey] = heartbeat
|
|
}
|
|
s.Len(workersByKey, 2)
|
|
|
|
worker1, exists := workersByKey[worker1Key]
|
|
s.True(exists, "worker1 should exist")
|
|
s.Equal(sharedTaskQueue, worker1.TaskQueue)
|
|
s.Equal(int32(1), worker1.TotalStickyCacheHit)
|
|
|
|
worker2, exists := workersByKey[worker2Key]
|
|
s.True(exists, "worker2 should exist")
|
|
s.Equal(sharedTaskQueue, worker2.TaskQueue)
|
|
s.Equal(int32(2), worker2.TotalStickyCacheHit)
|
|
|
|
// Verify new Workers field (WorkerListInfo)
|
|
listInfos := resp.GetWorkers()
|
|
s.Len(listInfos, 2)
|
|
listInfoByKey := make(map[string]*workerpb.WorkerListInfo)
|
|
for _, info := range listInfos {
|
|
listInfoByKey[info.GetWorkerInstanceKey()] = info
|
|
}
|
|
|
|
listInfo1, exists := listInfoByKey[worker1Key]
|
|
s.True(exists, "worker1 list info should exist")
|
|
s.Equal(sharedTaskQueue, listInfo1.GetTaskQueue())
|
|
|
|
listInfo2, exists := listInfoByKey[worker2Key]
|
|
s.True(exists, "worker2 list info should exist")
|
|
s.Equal(sharedTaskQueue, listInfo2.GetTaskQueue())
|
|
}
|
|
{
|
|
nonExistentWorkerKey := s.tv.WorkerIdentity() + "_nonexistent"
|
|
resp, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: fmt.Sprintf("WorkerInstanceKey='%s'", nonExistentWorkerKey),
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(resp)
|
|
s.Empty(resp.GetWorkersInfo()) //nolint:staticcheck // testing deprecated field
|
|
s.Empty(resp.GetWorkers())
|
|
}
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) TestWorkerRegistry_ListWorkersPagination() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
sharedTaskQueue := s.tv.TaskQueue().Name
|
|
|
|
// Create 5 workers with predictable keys for pagination testing
|
|
workerKeys := make([]string, 5)
|
|
heartbeats := make([]*workerpb.WorkerHeartbeat, 5)
|
|
for i := range 5 {
|
|
workerKeys[i] = fmt.Sprintf("%s_worker_%02d", s.tv.WorkerIdentity(), i)
|
|
heartbeats[i] = &workerpb.WorkerHeartbeat{
|
|
WorkerInstanceKey: workerKeys[i],
|
|
TaskQueue: sharedTaskQueue,
|
|
TotalStickyCacheHit: int32(i),
|
|
}
|
|
}
|
|
|
|
// Record all worker heartbeats
|
|
hbResp, err := s.FrontendClient().RecordWorkerHeartbeat(ctx, &workflowservice.RecordWorkerHeartbeatRequest{
|
|
Namespace: s.Namespace().String(),
|
|
WorkerHeartbeat: heartbeats,
|
|
})
|
|
s.NoError(err)
|
|
s.NotNil(hbResp)
|
|
|
|
query := fmt.Sprintf("TaskQueue='%s'", sharedTaskQueue)
|
|
|
|
// Page 1: Request first 2 workers
|
|
resp1, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: query,
|
|
PageSize: 2,
|
|
})
|
|
s.NoError(err)
|
|
s.Len(resp1.GetWorkersInfo(), 2)
|
|
s.NotEmpty(resp1.GetNextPageToken(), "should have next page token")
|
|
s.Equal(workerKeys[0], resp1.GetWorkersInfo()[0].GetWorkerHeartbeat().GetWorkerInstanceKey())
|
|
s.Equal(workerKeys[1], resp1.GetWorkersInfo()[1].GetWorkerHeartbeat().GetWorkerInstanceKey())
|
|
|
|
// Page 2: Request next 2 workers using token from page 1
|
|
resp2, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: query,
|
|
PageSize: 2,
|
|
NextPageToken: resp1.GetNextPageToken(),
|
|
})
|
|
s.NoError(err)
|
|
s.Len(resp2.GetWorkersInfo(), 2)
|
|
s.NotEmpty(resp2.GetNextPageToken(), "should have next page token")
|
|
s.Equal(workerKeys[2], resp2.GetWorkersInfo()[0].GetWorkerHeartbeat().GetWorkerInstanceKey())
|
|
s.Equal(workerKeys[3], resp2.GetWorkersInfo()[1].GetWorkerHeartbeat().GetWorkerInstanceKey())
|
|
|
|
// Page 3: Request remaining workers using token from page 2
|
|
resp3, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: query,
|
|
PageSize: 2,
|
|
NextPageToken: resp2.GetNextPageToken(),
|
|
})
|
|
s.NoError(err)
|
|
s.Len(resp3.GetWorkersInfo(), 1, "last page should have 1 worker")
|
|
s.Empty(resp3.GetNextPageToken(), "should not have next page token on last page")
|
|
s.Equal(workerKeys[4], resp3.GetWorkersInfo()[0].GetWorkerHeartbeat().GetWorkerInstanceKey())
|
|
}
|
|
|
|
func (s *WorkerRegistryTestSuite) TestWorkerRegistry_SendHeartbeatViaPollNexusTask() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
nexusWorkerKey := s.tv.WorkerIdentity() + "_nexus"
|
|
nexusTaskQueue := s.tv.TaskQueue().Name
|
|
|
|
heartbeat := &workerpb.WorkerHeartbeat{
|
|
WorkerInstanceKey: nexusWorkerKey,
|
|
TaskQueue: nexusTaskQueue,
|
|
TotalStickyCacheHit: 3,
|
|
}
|
|
|
|
// Send worker heartbeat via PollNexusTaskQueue in a goroutine.
|
|
// This is because PollNexusTaskQueue is a blocking call and we need to wait for the heartbeat to be registered.
|
|
go func() {
|
|
_, _ = s.FrontendClient().PollNexusTaskQueue(ctx, &workflowservice.PollNexusTaskQueueRequest{
|
|
Namespace: s.Namespace().String(),
|
|
TaskQueue: &taskqueuepb.TaskQueue{Name: nexusTaskQueue},
|
|
Identity: nexusWorkerKey,
|
|
WorkerHeartbeat: []*workerpb.WorkerHeartbeat{heartbeat},
|
|
})
|
|
}()
|
|
|
|
// Verify heartbeat was registered.
|
|
s.EventuallyWithT(func(t *assert.CollectT) {
|
|
resp, err := s.FrontendClient().ListWorkers(ctx, &workflowservice.ListWorkersRequest{
|
|
Namespace: s.Namespace().String(),
|
|
Query: fmt.Sprintf("WorkerInstanceKey='%s'", nexusWorkerKey),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
workers := resp.GetWorkersInfo()
|
|
require.Len(t, workers, 1)
|
|
|
|
workerHeartbeat := workers[0].GetWorkerHeartbeat()
|
|
require.Equal(t, heartbeat.WorkerInstanceKey, workerHeartbeat.WorkerInstanceKey)
|
|
require.Equal(t, heartbeat.TaskQueue, workerHeartbeat.TaskQueue)
|
|
require.Equal(t, heartbeat.TotalStickyCacheHit, workerHeartbeat.TotalStickyCacheHit)
|
|
}, 2*time.Minute, 100*time.Millisecond, "Worker heartbeat should be registered via PollNexusTaskQueue")
|
|
}
|