agent-framework-go/workflow/protocol_test.go at main · microsoft/agent-framework-go · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
// Copyright (c) Microsoft. All rights reserved.

package workflow

import (
	"errors"
	"reflect"
	"slices"
	"testing"
)

type (
	protocolBuilderInput  struct{}
	protocolBuilderOutput struct{}
	protocolBuilderSend   struct{}
	protocolBuilderYield  struct{}
)

func TestProtocolBuilderBuildIncludesDeclaredAndAutomaticTypes(t *testing.T) {
	inputType := reflect.TypeFor[protocolBuilderInput]()
	outputType := reflect.TypeFor[protocolBuilderOutput]()
	explicitSend := reflect.TypeFor[protocolBuilderSend]()
	explicitYield := reflect.TypeFor[protocolBuilderYield]()

	protocol, err := newProtocolBuilderWithHandler(inputType, outputType).
		SendsMessageType(explicitSend, nil, explicitSend).
		YieldsOutputType(explicitYield, nil, explicitYield).
		build(&Executor{})
	if err != nil {
		t.Fatalf("build error = %v", err)
	}

	if !containsReflectType(protocol.describe().Accepts, inputType) {
		t.Fatalf("Accepts = %v, want %v", protocol.describe().Accepts, inputType)
	}
	if got := protocol.describe().Sends; !reflect.DeepEqual(got, []reflect.Type{explicitSend, outputType}) {
		t.Fatalf("sends = %v, want [%v %v]", got, explicitSend, outputType)
	}
	if got := protocol.describe().Yields; !reflect.DeepEqual(got, []reflect.Type{explicitYield, outputType}) {
		t.Fatalf("yields = %v, want [%v %v]", got, explicitYield, outputType)
	}
}

func TestExecutorProtocolDescribeReturnsCachedDescriptorValue(t *testing.T) {
	protocol, err := newProtocolBuilderWithHandler(reflect.TypeFor[protocolBuilderInput](), reflect.TypeFor[protocolBuilderOutput]()).
		SendsMessageType(reflect.TypeFor[protocolBuilderSend]()).
		YieldsOutputType(reflect.TypeFor[protocolBuilderYield]()).
		build(&Executor{})
	if err != nil {
		t.Fatalf("build error = %v", err)
	}

	first := protocol.describe()
	second := protocol.describe()
	if !reflect.DeepEqual(first, second) {
		t.Fatalf("describe returned different descriptors: %+v != %+v", first, second)
	}
	if allocations := testing.AllocsPerRun(1000, func() { _ = protocol.describe() }); allocations != 0 {
		t.Fatalf("describe allocations = %v, want 0", allocations)
	}
}

func TestProtocolBuilderBuildRespectsAutoReturnOptions(t *testing.T) {
	inputType := reflect.TypeFor[protocolBuilderInput]()
	outputType := reflect.TypeFor[protocolBuilderOutput]()

	protocol, err := newProtocolBuilderWithHandler(inputType, outputType).build(&Executor{
		AutoSendMessageHandlerResultObject: new(false),
		AutoYieldOutputHandlerResultObject: new(false),
	})
	if err != nil {
		t.Fatalf("build error = %v", err)
	}
	if containsReflectType(protocol.describe().Sends, outputType) {
		t.Fatalf("sends = %v, want no automatic output type %v", protocol.describe().Sends, outputType)
	}
	if containsReflectType(protocol.describe().Yields, outputType) {
		t.Fatalf("yields = %v, want no automatic output type %v", protocol.describe().Yields, outputType)
	}
}

func TestProtocolBuilderConfigureRoutesErrorReturnsFromBuild(t *testing.T) {
	wantErr := errors.New("route setup failed")
	var pb ProtocolBuilder
	_, err := pb.ConfigureRoutes(func(*RouteBuilder) (*RouteBuilder, error) {
		return nil, wantErr
	}).build(&Executor{})
	if !errors.Is(err, wantErr) {
		t.Fatalf("build error = %v, want %v", err, wantErr)
	}
}

func newProtocolBuilderWithHandler(inputType, outputType reflect.Type) *ProtocolBuilder {
	var pb ProtocolBuilder
	pb.RouteBuilder.AddHandlerRaw(inputType, outputType, func(*Context, any) (any, error) {
		return protocolBuilderOutput{}, nil
	})
	return &pb
}

func containsReflectType(types []reflect.Type, want reflect.Type) bool {
	return slices.Contains(types, want)
}