forked from sashabaranov/go-openai
-
Notifications
You must be signed in to change notification settings - Fork 0
/
helpers.go
53 lines (44 loc) · 1.66 KB
/
helpers.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
package test
import (
"github.com/ceerdecy/go-openai/internal/test/checks"
"net/http"
"os"
"testing"
)
// CreateTestFile creates a fake file with "hello" as the content.
func CreateTestFile(t *testing.T, path string) {
file, err := os.Create(path)
checks.NoError(t, err, "failed to create file")
if _, err = file.WriteString("hello"); err != nil {
t.Fatalf("failed to write to file %v", err)
}
file.Close()
}
// CreateTestDirectory creates a temporary folder which will be deleted when cleanup is called.
func CreateTestDirectory(t *testing.T) (path string, cleanup func()) {
t.Helper()
path, err := os.MkdirTemp(os.TempDir(), "")
checks.NoError(t, err)
return path, func() { os.RemoveAll(path) }
}
// TokenRoundTripper is a struct that implements the RoundTripper
// interface, specifically to handle the authentication token by adding a token
// to the request header. We need this because the API requires that each
// request include a valid API token in the headers for authentication and
// authorization.
type TokenRoundTripper struct {
Token string
Fallback http.RoundTripper
}
// RoundTrip takes an *http.Request as input and returns an
// *http.Response and an error.
//
// It is expected to use the provided request to create a connection to an HTTP
// server and return the response, or an error if one occurred. The returned
// Response should have its Body closed. If the RoundTrip method returns an
// error, the Client's Get, Head, Post, and PostForm methods return the same
// error.
func (t *TokenRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
req.Header.Set("Authorization", "Bearer "+t.Token)
return t.Fallback.RoundTrip(req)
}