-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathsync_test.go
More file actions
70 lines (50 loc) · 1.96 KB
/
sync_test.go
File metadata and controls
70 lines (50 loc) · 1.96 KB
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
package llmberjack
import (
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)
func TestSyncAll(t *testing.T) {
e := errors.New("error")
p1 := NewMockProvider()
p1.On("Init", mock.Anything).Return(nil)
llm, _ := New(WithDefaultProvider(p1))
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(MockMessage{"first"}, nil).After(100 * time.Millisecond).Once()
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(nil, e).After(300 * time.Millisecond).Once()
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(MockMessage{"third"}, nil).After(500 * time.Millisecond).Once()
responses := All(t.Context(), llm,
NewUntypedRequest(),
NewUntypedRequest(),
NewUntypedRequest())
assert.Len(t, responses, 3)
values := lo.Map(responses, func(resp AsyncResponse[string], _ int) any {
if resp.Error != nil {
return resp.Error
}
output, err := resp.Response.Get(0)
assert.Nil(t, err)
return output
})
assert.ElementsMatch(t, values, []any{"first", e, "third"})
}
func TestSyncRace(t *testing.T) {
p1 := NewMockProvider()
p1.On("Init", mock.Anything).Return(nil)
llm, _ := New(WithDefaultProvider(p1))
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(nil, errors.New("error")).Once()
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(MockMessage{"first"}, nil).After(200 * time.Millisecond).Once()
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(MockMessage{"second"}, nil).After(400 * time.Millisecond).Once()
p1.On("ChatCompletion", mock.Anything, llm, mock.Anything).Return(MockMessage{"third"}, nil).After(500 * time.Millisecond).Once()
resp, err := Race(t.Context(), llm,
NewUntypedRequest(),
NewUntypedRequest(),
NewUntypedRequest())
assert.Nil(t, err)
assert.NotNil(t, resp)
output, err := resp.Get(0)
assert.Nil(t, err)
assert.Equal(t, "first", output)
}