Testing Guide¶

import (
    "testing"
    "time"

    helixtesting "github.com/arvenza/helix/pkg/testing"
)

func TestMyRoute(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    // Add your route to the context.
    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "my-workflow",
        FromURI: "direct://start",
        Processors: []engine.Processor{
            engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                body := ex.In().BodyAsString()
                ex.In().SetBody("processed: " + body)
                return nil
            }),
            // Send to a mock endpoint for verification.
            mustProducer(ts.Context, "mock://result"),
        },
    })

    // Start the context (starts all workflows).
    ts.Start(t)
    defer ts.Stop()

    // Set expectations on the mock endpoint.
    result := ts.GetMockEndpoint("result")
    result.ExpectedMessageCount(1)

    // Send a test message.
    if err := ts.Producer.SendBody("direct://start", "hello"); err != nil {
        t.Fatal(err)
    }

    // Assert expectations.
    if err := result.AssertIsSatisfied(5 * time.Second); err != nil {
        t.Fatal(err)
    }
}

// Send a body to an endpoint.
err := ts.Producer.SendBody("direct://start", "hello")

// Send a body with headers.
err := ts.Producer.SendBodyAndHeaders("direct://start", "hello", map[string]any{
    "Content-Type": "application/json",
    "X-Correlation": "abc-123",
})

// Send a request and receive a response body.
result, err := ts.Producer.RequestBody("direct://start", "ping")
// result contains the Out body (or In body if no Out is set).

// Send with headers and receive a response.
result, err := ts.Producer.RequestBodyAndHeaders("direct://start", requestData, headers)

ts := helixtesting.NewTestSupport()
mockEp := ts.GetMockEndpoint("result")

comp := mock.New()
ep, err := comp.CreateEndpoint("mock://result", "result", nil)
mockEp := ep.(*mock.Endpoint)

// Expect exactly N messages.
mockEp.ExpectedMessageCount(3)

// Expect at least N messages.
mockEp.SetMinimumExpectedMessageCount(2)

// Expect specific bodies in exact order.
mockEp.ExpectedBodiesReceived("alpha", "beta", "gamma")

// Expect specific bodies in any order.
mockEp.ExpectedBodiesReceivedInAnyOrder("gamma", "alpha", "beta")

// Every received message must have this header value.
mockEp.ExpectedHeaderReceived("X-Source", "test")

// Expect specific header values across messages, in any order.
mockEp.ExpectedHeaderValuesReceivedInAnyOrder("X-ID", "c", "a", "b")

mockEp.Expects(func(index int, ex *exchange.Exchange) error {
    body := ex.In().BodyAsString()
    if len(body) == 0 {
        return fmt.Errorf("expected non-empty body at index %d", index)
    }
    return nil
})

// Verify ascending order.
mockEp.ExpectsAscending(myExpression)

// Verify descending order.
mockEp.ExpectsDescending(myExpression)

// Verify no duplicate values.
mockEp.ExpectsNoDuplicates(myExpression)

// Set a reply body on every received exchange.
mockEp.ReturnReplyBody("reply-data")

// Set reply headers.
mockEp.ReturnReplyHeader("X-Reply", "yes")

mockEp.WhenAnyExchangeReceived(engine.ProcessorFunc(
    func(ctx context.Context, ex *exchange.Exchange) error {
        ex.In().SetBody("transformed-" + ex.In().BodyAsString())
        return nil
    },
))

// Wait up to the given duration for all expectations to be satisfied.
err := mockEp.AssertIsSatisfied(5 * time.Second)

// Override the default wait time (5 seconds).
mockEp.SetResultWaitTime(10 * time.Second)

// Query received state.
count := mockEp.ReceivedCount()
exchanges := mockEp.Received() // returns []*exchange.Exchange

// Reset for reuse between test phases.
mockEp.Reset()

ts := helixtesting.NewTestSupport()

ts.Context.AddRoute(productionRoute) // FromURI: "kafka://orders"

// Replace Kafka with direct:// for testing.
advice := ts.AdviceWith("my-workflow")
err := advice.ReplaceFromWith("direct://test-input")

// Replace any endpoint matching "log" or "http" with a mock.
err := advice.MockEndpoints("log", "http")

err := advice.InterceptSendToEndpoint("log", engine.ProcessorFunc(
    func(ctx context.Context, ex *exchange.Exchange) error {
        // Inspect or modify before the original send.
        log.Printf("intercepted: %v", ex.In().Body())
        return nil
    },
))

// Add a processor at the beginning.
advice.WeaveAddFirst(myValidator)

// Add a processor at the end.
advice.WeaveAddLast(myAuditor)

// Insert before/after/replace a specific step by ID.
advice.WeaveByID("validate-step").Before(myPreProcessor)
advice.WeaveByID("validate-step").After(myPostProcessor)
advice.WeaveByID("validate-step").Replace(myNewValidator)
advice.WeaveByID("validate-step").Remove()

tagged := helixtesting.NewTaggedProcessor("validate-step",
    engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
        // validation logic
        return nil
    }),
)

workflow := &engine.Workflow{
    ID:         "my-workflow",
    FromURI:    "direct://start",
    Processors: []engine.Processor{tagged, otherProc},
}

func TestOrderProcessing(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "order-route",
        FromURI: "direct://orders",
        Processors: []engine.Processor{
            helixtesting.NewTaggedProcessor("validate",
                engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                    // Validate order
                    return nil
                }),
            ),
            helixtesting.NewTaggedProcessor("enrich",
                engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                    ex.In().SetHeader("enriched", true)
                    return nil
                }),
            ),
            mustProducer(ts.Context, "mock://result"),
        },
    })

    ts.Start(t)
    defer ts.Stop()

    result := ts.GetMockEndpoint("result")
    result.ExpectedMessageCount(1)
    result.ExpectedHeaderReceived("enriched", true)

    ts.Producer.SendBody("direct://orders", `{"item":"widget","qty":5}`)

    if err := result.AssertIsSatisfied(5 * time.Second); err != nil {
        t.Fatal(err)
    }
}

func TestAsyncPipeline(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    // Add workflows that use seda:// (async).
    ts.Context.AddRoute(ingestRoute)   // direct://input -> seda://process
    ts.Context.AddRoute(processRoute)  // seda://process -> mock://output

    ts.Start(t)
    defer ts.Stop()

    // Wait for 3 exchanges to complete on the "process-route".
    notify := helixtesting.NewNotifyBuilder(ts.Context).
        FromRoute("process-route").
        WhenDone(3).
        Create()

    // Send 3 messages.
    for i := 0; i < 3; i++ {
        ts.Producer.SendBody("direct://input", fmt.Sprintf("msg-%d", i))
    }

    // Block until 3 exchanges are done (or timeout).
    if err := notify.Matches(10 * time.Second); err != nil {
        t.Fatal(err)
    }
}

notify := helixtesting.NewNotifyBuilder(ts.Context).
    FromRoute("my-workflow").
    WhenDone(1).
    Filter(func(ctx context.Context, ex *exchange.Exchange) (bool, error) {
        // Only count exchanges with a specific header.
        return ex.In().Header("X-Priority") == "high", nil
    }).
    Create()

func TestWithCoverage(t *testing.T) {
    ts := helixtesting.NewTestSupport()
    coverage := helixtesting.NewRouteCoverage()

    // Register workflows and their processor IDs.
    coverage.RegisterRoute("my-workflow", []string{"validate", "transform", "send"})

    // Attach as an interceptor so hits are recorded automatically.
    ts.Context.AddInterceptor(coverage)

    // ... add workflows, start, run tests ...

    // After tests, inspect coverage.
    report := coverage.Report()
    t.Logf("Coverage: %.1f%% (%d/%d processors)",
        report.Coverage, report.CoveredProcessors, report.TotalProcessors)

    // Print a formatted report.
    coverage.PrintReport(os.Stdout)

    // Fail if coverage is below threshold.
    if report.Coverage < 80.0 {
        t.Errorf("route coverage %.1f%% is below 80%% threshold", report.Coverage)
    }
}

ts := helixtesting.NewTestSupport()

// Production route consumes from Kafka.
ts.Context.AddWorkflow(&engine.Workflow{
    ID:      "ingest",
    FromURI: "kafka://orders?group=processor",
    Processors: []engine.Processor{orderProcessor, mustProducer(ts.Context, "mock://output")},
})

// Replace Kafka with direct:// for testing.
advice := ts.AdviceWith("ingest")
advice.ReplaceFromWith("direct://test-orders")

ts.Start(t)
defer ts.Stop()

// Now send directly.
ts.Producer.SendBody("direct://test-orders", testOrderJSON)

func TestRequestReply(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "echo-route",
        FromURI: "direct://echo",
        Processors: []engine.Processor{
            engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                out := exchange.NewMessage()
                out.SetBody("echo: " + ex.In().BodyAsString())
                ex.SetOut(out)
                return nil
            }),
        },
    })

    ts.Start(t)
    defer ts.Stop()

    result, err := ts.Producer.RequestBody("direct://echo", "ping")
    if err != nil {
        t.Fatal(err)
    }
    if result != "echo: ping" {
        t.Errorf("got %v, want %q", result, "echo: ping")
    }
}

func TestAsyncTransform(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    // Workflow 1: direct -> seda (enqueue).
    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "enqueue",
        FromURI: "direct://input",
        Processors: []engine.Processor{mustProducer(ts.Context, "seda://work")},
    })

    // Workflow 2: seda -> process -> mock (dequeue and process).
    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "dequeue",
        FromURI: "seda://work",
        Processors: []engine.Processor{
            engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                ex.In().SetBody("done:" + ex.In().BodyAsString())
                return nil
            }),
            mustProducer(ts.Context, "mock://output"),
        },
    })

    ts.Start(t)
    defer ts.Stop()

    output := ts.GetMockEndpoint("output")
    output.ExpectedBodiesReceived("done:task-1")

    ts.Producer.SendBody("direct://input", "task-1")

    if err := output.AssertIsSatisfied(5 * time.Second); err != nil {
        t.Fatal(err)
    }
}

// pkg/agent/llm/model.go
type ChatModel interface {
    Generate(ctx context.Context, messages []*Message, opts ...Option) (*Message, error)
    Stream(ctx context.Context, messages []*Message, opts ...Option) (*StreamReader, error)
    BindTools(tools []ToolDefinition) error
    Provider() string
    Model() string
}

package mytest

import (
    "context"
    "github.com/arvenza/helix/pkg/agent/llm"
)

// MockChatModel returns predetermined responses for testing.
type MockChatModel struct {
    GenerateFn func(ctx context.Context, messages []*llm.Message, opts ...llm.Option) (*llm.Message, error)
    StreamFn   func(ctx context.Context, messages []*llm.Message, opts ...llm.Option) (*llm.StreamReader, error)
    BoundTools []llm.ToolDefinition
}

func (m *MockChatModel) Generate(ctx context.Context, messages []*llm.Message, opts ...llm.Option) (*llm.Message, error) {
    if m.GenerateFn != nil {
        return m.GenerateFn(ctx, messages, opts...)
    }
    return &llm.Message{
        Role:    llm.RoleAssistant,
        Content: "mock response",
    }, nil
}

func (m *MockChatModel) Stream(ctx context.Context, messages []*llm.Message, opts ...llm.Option) (*llm.StreamReader, error) {
    if m.StreamFn != nil {
        return m.StreamFn(ctx, messages, opts...)
    }
    ch := make(chan llm.StreamChunk, 1)
    ch <- llm.StreamChunk{Content: "mock stream", Done: true}
    close(ch)
    return llm.NewStreamReader(ch), nil
}

func (m *MockChatModel) BindTools(tools []llm.ToolDefinition) error {
    m.BoundTools = tools
    return nil
}

func (m *MockChatModel) Provider() string { return "mock" }
func (m *MockChatModel) Model() string    { return "mock-model" }

func TestAgentRoute(t *testing.T) {
    mockLLM := &MockChatModel{
        GenerateFn: func(ctx context.Context, msgs []*llm.Message, opts ...llm.Option) (*llm.Message, error) {
            // Inspect the prompt.
            lastMsg := msgs[len(msgs)-1]
            if lastMsg.Content == "What is 2+2?" {
                return &llm.Message{
                    Role:    llm.RoleAssistant,
                    Content: "4",
                }, nil
            }
            return &llm.Message{
                Role:    llm.RoleAssistant,
                Content: "I don't know",
            }, nil
        },
    }

    // Use mockLLM wherever your agent code accepts an llm.ChatModel.
    // For example, pass it into an agent constructor or set it on the context.
    _ = mockLLM
}

func TestAgentToolCall(t *testing.T) {
    callCount := 0
    mockLLM := &MockChatModel{
        GenerateFn: func(ctx context.Context, msgs []*llm.Message, opts ...llm.Option) (*llm.Message, error) {
            callCount++
            if callCount == 1 {
                // First call: model requests a tool call.
                return &llm.Message{
                    Role: llm.RoleAssistant,
                    ToolCalls: []llm.ToolCall{{
                        ID:   "call-1",
                        Type: "function",
                        Function: struct {
                            Name      string `json:"name"`
                            Arguments string `json:"arguments"`
                        }{
                            Name:      "get_weather",
                            Arguments: `{"city":"London"}`,
                        },
                    }},
                }, nil
            }
            // Second call: model produces final answer after receiving tool result.
            return &llm.Message{
                Role:    llm.RoleAssistant,
                Content: "The weather in London is sunny.",
            }, nil
        },
    }

    _ = mockLLM // wire into your agent under test
}

func TestStreamingAgent(t *testing.T) {
    mockLLM := &MockChatModel{
        StreamFn: func(ctx context.Context, msgs []*llm.Message, opts ...llm.Option) (*llm.StreamReader, error) {
            ch := make(chan llm.StreamChunk, 3)
            go func() {
                defer close(ch)
                ch <- llm.StreamChunk{Content: "Hello "}
                ch <- llm.StreamChunk{Content: "world"}
                ch <- llm.StreamChunk{Content: "!", Done: true}
            }()
            return llm.NewStreamReader(ch), nil
        },
    }

    // Collect the full streamed response.
    reader, err := mockLLM.Stream(context.Background(), []*llm.Message{
        {Role: llm.RoleUser, Content: "Say hello"},
    })
    if err != nil {
        t.Fatal(err)
    }
    result, err := reader.Collect()
    if err != nil {
        t.Fatal(err)
    }
    if result != "Hello world!" {
        t.Errorf("got %q, want %q", result, "Hello world!")
    }
}

func TestWithMockProvider(t *testing.T) {
    llm.RegisterProvider("mock", func(config llm.ModelConfig) (llm.ChatModel, error) {
        return &MockChatModel{}, nil
    })

    model, err := llm.NewChatModel(llm.ModelConfig{
        Provider: "mock",
        Model:    "test",
    })
    if err != nil {
        t.Fatal(err)
    }
    // Use model in your agent test.
    _ = model
}

func TestErrorHandling(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "error-route",
        FromURI: "direct://input",
        Processors: []engine.Processor{
            engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                return fmt.Errorf("simulated failure")
            }),
        },
    })

    ts.Start(t)
    defer ts.Stop()

    // Send a message that will trigger the error.
    err := ts.Producer.SendBody("direct://input", "bad-data")
    if err == nil {
        t.Fatal("expected error from failed route")
    }
    if !strings.Contains(err.Error(), "simulated failure") {
        t.Errorf("unexpected error message: %v", err)
    }
}

func TestDeadLetterChannel(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    dlq := ts.GetMockEndpoint("dead-letter")
    dlq.ExpectedMessageCount(1)

    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "dlq-route",
        FromURI: "direct://input",
        Processors: []engine.Processor{
            engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                return fmt.Errorf("permanent failure")
            }),
        },
        // Configure ErrorHandler with DeadLetterChannel: "mock://dead-letter"
    })

    ts.Start(t)
    defer ts.Stop()

    _ = ts.Producer.SendBody("direct://input", "will-fail")

    if err := dlq.AssertIsSatisfied(5 * time.Second); err != nil {
        t.Fatal(err)
    }

    // Inspect the dead-lettered exchange.
    received := dlq.Received()
    if len(received) != 1 {
        t.Fatalf("expected 1 dead-lettered exchange, got %d", len(received))
    }
}

func TestExchangeErrorState(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    output := ts.GetMockEndpoint("output")
    output.ExpectedMessageCount(1)
    output.Expects(func(index int, ex *exchange.Exchange) error {
        if ex.Exception() == nil {
            return fmt.Errorf("expected exception on exchange at index %d", index)
        }
        return nil
    })

    // Set up a workflow that catches errors and forwards them to mock://output.
    // ...
}

func TestProcessorError(t *testing.T) {
    validator := engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
        body := ex.In().BodyAsString()
        if body == "" {
            return fmt.Errorf("validation failed: body is empty")
        }
        return nil
    })

    // Test the error case.
    ex := exchange.New(context.Background())
    ex.In().SetBody("")

    err := validator.Process(context.Background(), ex)
    if err == nil {
        t.Fatal("expected validation error")
    }
    if !strings.Contains(err.Error(), "validation failed") {
        t.Errorf("unexpected error: %v", err)
    }

    // Test the success case.
    ex2 := exchange.New(context.Background())
    ex2.In().SetBody("valid data")

    err = validator.Process(context.Background(), ex2)
    if err != nil {
        t.Fatalf("unexpected error: %v", err)
    }
}

func TestRetryBehavior(t *testing.T) {
    ts := helixtesting.NewTestSupport()

    attempts := 0
    ts.Context.AddWorkflow(&engine.Workflow{
        ID:      "retry-route",
        FromURI: "direct://input",
        Processors: []engine.Processor{
            engine.ProcessorFunc(func(ctx context.Context, ex *exchange.Exchange) error {
                attempts++
                if attempts < 3 {
                    return fmt.Errorf("transient error (attempt %d)", attempts)
                }
                return nil
            }),
            mustProducer(ts.Context, "mock://result"),
        },
        // Configure ErrorHandler with MaxRedeliveries: 3
    })

    ts.Start(t)
    defer ts.Stop()

    result := ts.GetMockEndpoint("result")
    result.ExpectedMessageCount(1)

    ts.Producer.SendBody("direct://input", "data")

    if err := result.AssertIsSatisfied(5 * time.Second); err != nil {
        t.Fatal(err)
    }
    if attempts < 3 {
        t.Errorf("expected at least 3 attempts, got %d", attempts)
    }
}

// Test a tool via its auto-generated URI
result, err := ts.Producer.RequestBody("tool://search", map[string]any{"query": "test"})
if err != nil {
    t.Fatal(err)
}

// Test an agent route directly (with a mock LLM wired in)
resp, err := agent.Run(ctx, "What is 2+2?")
if err != nil {
    t.Fatal(err)
}
if resp != "4" {
    t.Errorf("unexpected response: %s", resp)
}

ts.Start(t)
defer ts.Stop()

mockEp.SetResultWaitTime(5 * time.Second)
err := mockEp.AssertIsSatisfied(5 * time.Second)

mockEp.Reset()
mockEp.ExpectedMessageCount(2)
// ... send more messages ...

func TestRouteA(t *testing.T) {
    t.Parallel()
    ts := helixtesting.NewTestSupport()
    // ...
}

func TestRouteB(t *testing.T) {
    t.Parallel()
    ts := helixtesting.NewTestSupport()
    // ...
}