43 lines
959 B
Go
43 lines
959 B
Go
|
package trace_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/cloudwego/eino/compose"
|
||
|
|
"github.com/hhs/camtalk/internal/trace"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestEinoContextPropagation(t *testing.T) {
|
||
|
|
ctx := context.Background()
|
||
|
|
testTraceID := "01J5TEST123456789"
|
||
|
|
ctx = trace.WithTraceID(ctx, testTraceID)
|
||
|
|
|
||
|
|
var capturedTraceID string
|
||
|
|
|
||
|
|
g := compose.NewGraph[string, string]()
|
||
|
|
g.AddLambdaNode("test_node", compose.InvokableLambda(
|
||
|
|
func(ctx context.Context, input string) (string, error) {
|
||
|
|
capturedTraceID = trace.GetTraceID(ctx)
|
||
|
|
return "ok", nil
|
||
|
|
},
|
||
|
|
))
|
||
|
|
g.AddEdge(compose.START, "test_node")
|
||
|
|
g.AddEdge("test_node", compose.END)
|
||
|
|
|
||
|
|
runnable, err := g.Compile(ctx)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("compile failed: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
_, err = runnable.Invoke(ctx, "test_input")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("invoke failed: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
if capturedTraceID != testTraceID {
|
||
|
|
t.Errorf("trace_id lost in Eino propagation: got %q, want %q",
|
||
|
|
capturedTraceID, testTraceID)
|
||
|
|
}
|
||
|
|
}
|