-
Notifications
You must be signed in to change notification settings - Fork 1
/
banana_test.go
65 lines (51 loc) · 1.33 KB
/
banana_test.go
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
package banana
import (
"context"
"encoding/json"
"errors"
"fmt"
"testing"
"time"
)
const (
testingAPIKey = "demo"
testingModelKey = "carrot"
testingCarrotImageURL = "https://demo-images-banana.s3.us-west-1.amazonaws.com/image2.jpg"
)
func TestRunCarrot(t *testing.T) {
type input struct {
ImageURL string `json:"imageURL"`
}
in := input{
ImageURL: testingCarrotImageURL,
}
bytesIn, _ := json.Marshal(in)
bananaOut, err := Run(testingAPIKey, testingModelKey, bytesIn)
if err != nil {
panic(err)
}
type output []struct {
Caption string `json:"caption"`
}
out := output{}
json.Unmarshal(bananaOut.ModelOutputs, &out)
fmt.Println(out[0].Caption)
}
func TestContext(t *testing.T) {
buf, err := json.Marshal(map[string]string{
"imageURL": testingCarrotImageURL,
})
if err != nil {
t.Fatal(err)
}
// We put a stupidly low timeout, to make sure we get a timeout error.
// The laws of physics *should* prohibit a value of 1ms from ever being
// a valid timeout... we'd hope at least.
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
_, err = RunWithContext(ctx, testingAPIKey, testingModelKey, buf)
// Ensure that the context was cancelled.
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("expected context.DeadlineExceeded error, got %v", err)
}
}