mirror of
https://github.com/hanzoai/tasks.git
synced 2026-11-03 03:56:04 +00:00
nikki-dagandGitHub
139ad62c9d
## What changed?
1. Exposed metadata functionality:
- Made `SetContextMetadata()` public in MutableState interface
- Added constants for metadata keys (`MetadataKeyWorkflowType`,
`MetadataKeyWorkflowTaskQueue`)
- Added `ContextHasMetadata()` helper function to check if context has
metadata support
2. Fixed readonly operations:
- `QueryWorkflow`: Now explicitly calls `SetContextMetadata()` since
queries never close transactions (readonly operation)
- `RespondWorkflowTaskCompleted`: Calls `SetContextMetadata()` when
there are only messages (e.g., update rejections) and no commands, since
these also don't close transactions
3. Added tests for the new `ContextHasMetadata()` helper and
`SetContextMetadata()` functionality
## Why?
Context metadata (workflow type, task queue) is automatically populated
when mutable state transactions close. However, readonly operations
never close transactions, so the metadata was missing from their
contexts.
This change ensures that successful readonly operations (queries, update
rejections) also have workflow metadata populated in their contexts
## How did you test it?
- [ ] built
- [ ] run locally and tested manually
- [ ] covered by existing tests
- [x] added new unit test(s)
- [ ] added new functional test(s)
---
> [!NOTE]
> **Medium Risk**
> Touches core history API request paths and mutable state interfaces;
while behavior is additive (context tagging), incorrect invocation or
key usage could affect observability/metrics and tests across multiple
workflows.
>
> **Overview**
> Ensures workflow context metadata (workflow type/task queue) is
consistently populated even for readonly paths that don’t close
mutable-state transactions.
>
> This exposes `MutableState.SetContextMetadata(ctx)` (renaming the
internal helper), standardizes metadata keys via
`contextutil.MetadataKeyWorkflowType`/`MetadataKeyWorkflowTaskQueue`,
adds `ContextHasMetadata()` for debugging, and explicitly calls
`SetContextMetadata` in `QueryWorkflow` and in
`RespondWorkflowTaskCompleted` when `Commands` is empty (e.g., update
rejections / heartbeats). Tests are expanded to assert metadata presence
across success and several error/edge cases.
>
> <sup>Written by [Cursor
Bugbot]() for commit
69c3140233241c2072e3bacff7580debe91ed884. This will update automatically
on new commits. Configure
[here]().</sup>
364 lines
10 KiB
Go
364 lines
10 KiB
Go
package contextutil
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestAddMetadataContext(t *testing.T) {
|
|
t.Run("adds metadata context to empty context", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctx = WithMetadataContext(ctx)
|
|
|
|
metadataCtx := getMetadataContext(ctx)
|
|
require.NotNil(t, metadataCtx)
|
|
assert.NotNil(t, metadataCtx.Metadata)
|
|
assert.Empty(t, metadataCtx.Metadata)
|
|
})
|
|
|
|
t.Run("returns new context with metadata", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctxWithMetadata := WithMetadataContext(ctx)
|
|
|
|
assert.NotEqual(t, ctx, ctxWithMetadata)
|
|
assert.Nil(t, getMetadataContext(ctx))
|
|
assert.NotNil(t, getMetadataContext(ctxWithMetadata))
|
|
})
|
|
}
|
|
|
|
func TestContextMetadataSet(t *testing.T) {
|
|
t.Run("sets string value successfully", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
success := ContextMetadataSet(ctx, "key1", "value1")
|
|
assert.True(t, success)
|
|
|
|
metadataCtx := getMetadataContext(ctx)
|
|
require.NotNil(t, metadataCtx)
|
|
assert.Equal(t, "value1", metadataCtx.Metadata["key1"])
|
|
})
|
|
|
|
t.Run("sets multiple values successfully", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
ContextMetadataSet(ctx, "key2", 42)
|
|
ContextMetadataSet(ctx, "key3", true)
|
|
|
|
metadataCtx := getMetadataContext(ctx)
|
|
require.NotNil(t, metadataCtx)
|
|
assert.Equal(t, "value1", metadataCtx.Metadata["key1"])
|
|
assert.Equal(t, 42, metadataCtx.Metadata["key2"])
|
|
assert.Equal(t, true, metadataCtx.Metadata["key3"])
|
|
})
|
|
|
|
t.Run("overwrites existing value", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
ContextMetadataSet(ctx, "key1", "value2")
|
|
|
|
metadataCtx := getMetadataContext(ctx)
|
|
require.NotNil(t, metadataCtx)
|
|
assert.Equal(t, "value2", metadataCtx.Metadata["key1"])
|
|
})
|
|
|
|
t.Run("returns false when context has no metadata", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
success := ContextMetadataSet(ctx, "key1", "value1")
|
|
assert.False(t, success)
|
|
})
|
|
|
|
t.Run("supports various value types", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
type customStruct struct {
|
|
Field string
|
|
}
|
|
|
|
testCases := []struct {
|
|
key string
|
|
value any
|
|
}{
|
|
{"string", "test"},
|
|
{"int", 123},
|
|
{"float", 3.14},
|
|
{"bool", true},
|
|
{"slice", []string{"a", "b", "c"}},
|
|
{"map", map[string]int{"a": 1, "b": 2}},
|
|
{"struct", customStruct{Field: "test"}},
|
|
{"nil", nil},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.key, func(t *testing.T) {
|
|
success := ContextMetadataSet(ctx, tc.key, tc.value)
|
|
assert.True(t, success)
|
|
|
|
metadataCtx := getMetadataContext(ctx)
|
|
require.NotNil(t, metadataCtx)
|
|
assert.Equal(t, tc.value, metadataCtx.Metadata[tc.key])
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestContextMetadataGet(t *testing.T) {
|
|
t.Run("retrieves existing value", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
|
|
value, ok := ContextMetadataGet(ctx, "key1")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "value1", value)
|
|
})
|
|
|
|
t.Run("returns false for non-existent key", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
value, ok := ContextMetadataGet(ctx, "nonexistent")
|
|
assert.False(t, ok)
|
|
assert.Nil(t, value)
|
|
})
|
|
|
|
t.Run("returns false when context has no metadata", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
value, ok := ContextMetadataGet(ctx, "key1")
|
|
assert.False(t, ok)
|
|
assert.Nil(t, value)
|
|
})
|
|
|
|
t.Run("retrieves nil value correctly", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(ctx, "nilKey", nil)
|
|
|
|
value, ok := ContextMetadataGet(ctx, "nilKey")
|
|
assert.True(t, ok)
|
|
assert.Nil(t, value)
|
|
})
|
|
|
|
t.Run("retrieves various value types", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
ContextMetadataSet(ctx, "string", "test")
|
|
ContextMetadataSet(ctx, "int", 42)
|
|
ContextMetadataSet(ctx, "slice", []int{1, 2, 3})
|
|
|
|
strVal, ok := ContextMetadataGet(ctx, "string")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "test", strVal)
|
|
|
|
intVal, ok := ContextMetadataGet(ctx, "int")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, 42, intVal)
|
|
|
|
sliceVal, ok := ContextMetadataGet(ctx, "slice")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, []int{1, 2, 3}, sliceVal)
|
|
})
|
|
}
|
|
|
|
func TestContextMetadataGetAll(t *testing.T) {
|
|
t.Run("retrieves all metadata", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
ContextMetadataSet(ctx, "key2", 42)
|
|
|
|
allMetadata := ContextMetadataGetAll(ctx)
|
|
assert.NotNil(t, allMetadata)
|
|
assert.Len(t, allMetadata, 2)
|
|
assert.Equal(t, "value1", allMetadata["key1"])
|
|
assert.Equal(t, 42, allMetadata["key2"])
|
|
})
|
|
|
|
t.Run("returns nil when context has no metadata", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
allMetadata := ContextMetadataGetAll(ctx)
|
|
assert.Nil(t, allMetadata)
|
|
})
|
|
|
|
t.Run("returns empty map when no metadata set", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
|
|
allMetadata := ContextMetadataGetAll(ctx)
|
|
assert.NotNil(t, allMetadata)
|
|
assert.Empty(t, allMetadata)
|
|
})
|
|
|
|
t.Run("returned map is a copy", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
|
|
allMetadata := ContextMetadataGetAll(ctx)
|
|
allMetadata["key2"] = "value2"
|
|
|
|
// Original should not be affected
|
|
_, ok := ContextMetadataGet(ctx, "key2")
|
|
assert.False(t, ok)
|
|
})
|
|
}
|
|
|
|
func TestGetMetadataContext(t *testing.T) {
|
|
t.Run("returns nil for context without metadata", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
metadataCtx := getMetadataContext(ctx)
|
|
assert.Nil(t, metadataCtx)
|
|
})
|
|
|
|
t.Run("returns metadata context when present", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
metadataCtx := getMetadataContext(ctx)
|
|
assert.NotNil(t, metadataCtx)
|
|
})
|
|
|
|
t.Run("returns nil for wrong type in context", func(t *testing.T) {
|
|
ctx := context.WithValue(context.Background(), metadataCtxKey, "wrong type")
|
|
metadataCtx := getMetadataContext(ctx)
|
|
assert.Nil(t, metadataCtx)
|
|
})
|
|
}
|
|
|
|
func TestMetadataContextWithContextCancellation(t *testing.T) {
|
|
t.Run("metadata survives context cancellation", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
ctx = WithMetadataContext(ctx)
|
|
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
|
|
// Cancel the context
|
|
cancel()
|
|
|
|
// Metadata should still be accessible
|
|
value, ok := ContextMetadataGet(ctx, "key1")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "value1", value)
|
|
|
|
// Should still be able to set new values
|
|
success := ContextMetadataSet(ctx, "key2", "value2")
|
|
assert.True(t, success)
|
|
|
|
value, ok = ContextMetadataGet(ctx, "key2")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "value2", value)
|
|
})
|
|
}
|
|
|
|
func TestMetadataContextIsolation(t *testing.T) {
|
|
t.Run("contexts with different metadata are isolated", func(t *testing.T) {
|
|
ctx1 := WithMetadataContext(context.Background())
|
|
ctx2 := WithMetadataContext(context.Background())
|
|
|
|
ContextMetadataSet(ctx1, "key", "value1")
|
|
ContextMetadataSet(ctx2, "key", "value2")
|
|
|
|
value1, ok1 := ContextMetadataGet(ctx1, "key")
|
|
value2, ok2 := ContextMetadataGet(ctx2, "key")
|
|
|
|
assert.True(t, ok1)
|
|
assert.True(t, ok2)
|
|
assert.Equal(t, "value1", value1)
|
|
assert.Equal(t, "value2", value2)
|
|
})
|
|
|
|
t.Run("child context does not inherit parent metadata", func(t *testing.T) {
|
|
parentCtx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(parentCtx, "key", "parent-value")
|
|
|
|
type testContextKey string
|
|
childCtx := context.WithValue(parentCtx, testContextKey("other-key"), "other-value")
|
|
|
|
// Child can still access parent's metadata context
|
|
value, ok := ContextMetadataGet(childCtx, "key")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "parent-value", value)
|
|
|
|
// Setting in child affects parent (same metadata context)
|
|
ContextMetadataSet(childCtx, "key2", "child-value")
|
|
value, ok = ContextMetadataGet(parentCtx, "key2")
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "child-value", value)
|
|
})
|
|
|
|
t.Run("adding metadata context twice creates new isolated context", func(t *testing.T) {
|
|
ctx1 := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(ctx1, "key", "value1")
|
|
|
|
ctx2 := WithMetadataContext(ctx1)
|
|
ContextMetadataSet(ctx2, "key", "value2")
|
|
|
|
value1, ok1 := ContextMetadataGet(ctx1, "key")
|
|
value2, ok2 := ContextMetadataGet(ctx2, "key")
|
|
|
|
assert.True(t, ok1)
|
|
assert.True(t, ok2)
|
|
assert.Equal(t, "value1", value1)
|
|
assert.Equal(t, "value2", value2)
|
|
})
|
|
}
|
|
|
|
func TestContextHasMetadata(t *testing.T) {
|
|
t.Run("returns true when context has metadata", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
assert.True(t, ContextHasMetadata(ctx))
|
|
})
|
|
|
|
t.Run("returns false for context without metadata", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
assert.False(t, ContextHasMetadata(ctx))
|
|
})
|
|
|
|
t.Run("returns true after setting metadata values", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(ctx, "key1", "value1")
|
|
ContextMetadataSet(ctx, "key2", "value2")
|
|
|
|
assert.True(t, ContextHasMetadata(ctx))
|
|
})
|
|
|
|
t.Run("returns true for empty metadata context", func(t *testing.T) {
|
|
ctx := WithMetadataContext(context.Background())
|
|
// No values set, but metadata context exists
|
|
assert.True(t, ContextHasMetadata(ctx))
|
|
})
|
|
|
|
t.Run("child context inherits metadata from parent", func(t *testing.T) {
|
|
parentCtx := WithMetadataContext(context.Background())
|
|
ContextMetadataSet(parentCtx, "key", "value")
|
|
|
|
type testContextKey string
|
|
childCtx := context.WithValue(parentCtx, testContextKey("other-key"), "other-value")
|
|
|
|
assert.True(t, ContextHasMetadata(parentCtx))
|
|
assert.True(t, ContextHasMetadata(childCtx))
|
|
})
|
|
|
|
t.Run("returns false for wrong type in context", func(t *testing.T) {
|
|
ctx := context.WithValue(context.Background(), metadataCtxKey, "wrong type")
|
|
assert.False(t, ContextHasMetadata(ctx))
|
|
})
|
|
|
|
t.Run("returns true for cancelled context with metadata", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
ctx = WithMetadataContext(ctx)
|
|
cancel()
|
|
|
|
assert.True(t, ContextHasMetadata(ctx))
|
|
})
|
|
|
|
t.Run("multiple contexts with metadata are independent", func(t *testing.T) {
|
|
ctx1 := WithMetadataContext(context.Background())
|
|
ctx2 := WithMetadataContext(context.Background())
|
|
ctx3 := context.Background()
|
|
|
|
assert.True(t, ContextHasMetadata(ctx1))
|
|
assert.True(t, ContextHasMetadata(ctx2))
|
|
assert.False(t, ContextHasMetadata(ctx3))
|
|
})
|
|
}
|