Skip to content

Commit

Permalink
Add files api tests (sashabaranov#238)
Browse files Browse the repository at this point in the history
* drop support for downloading files

* use form builder to submit files

* update doc

* add form builder tests
  • Loading branch information
sashabaranov committed Apr 8, 2023
1 parent 4dc1eda commit 33ca1da
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 54 deletions.
60 changes: 9 additions & 51 deletions files.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,7 @@ import (
"bytes"
"context"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/url"
"os"
)

Expand All @@ -33,77 +30,38 @@ type FilesList struct {
Files []File `json:"data"`
}

// isUrl is a helper function that determines whether the given FilePath
// is a remote URL or a local file path.
func isURL(path string) bool {
_, err := url.ParseRequestURI(path)
if err != nil {
return false
}

u, err := url.Parse(path)
if err != nil || u.Scheme == "" || u.Host == "" {
return false
}

return true
}

// CreateFile uploads a jsonl file to GPT3
// FilePath can be either a local file path or a URL.
// FilePath must be a local file path.
func (c *Client) CreateFile(ctx context.Context, request FileRequest) (file File, err error) {
var b bytes.Buffer
w := multipart.NewWriter(&b)

var fw io.Writer
builder := c.createFormBuilder(&b)

err = w.WriteField("purpose", request.Purpose)
err = builder.writeField("purpose", request.Purpose)
if err != nil {
return
}

fw, err = w.CreateFormFile("file", request.FileName)
fileData, err := os.Open(request.FilePath)
if err != nil {
return
}

var fileData io.ReadCloser
if isURL(request.FilePath) {
var remoteFile *http.Response
remoteFile, err = http.Get(request.FilePath)
if err != nil {
return
}

defer remoteFile.Body.Close()

// Check server response
if remoteFile.StatusCode != http.StatusOK {
err = fmt.Errorf("error, status code: %d, message: failed to fetch file", remoteFile.StatusCode)
return
}

fileData = remoteFile.Body
} else {
fileData, err = os.Open(request.FilePath)
if err != nil {
return
}
err = builder.createFormFile("file", fileData)
if err != nil {
return
}

_, err = io.Copy(fw, fileData)
err = builder.close()
if err != nil {
return
}

w.Close()

req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.fullURL("/files"), &b)
if err != nil {
return
}

req.Header.Set("Content-Type", w.FormDataContentType())
req.Header.Set("Content-Type", builder.formDataContentType())

err = c.sendRequest(req, &file)

Expand Down
54 changes: 51 additions & 3 deletions files_test.go
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
package openai_test
package openai //nolint:testpackage // testing private field

import (
. "github.com/sashabaranov/go-openai"
"github.com/sashabaranov/go-openai/internal/test"
"github.com/sashabaranov/go-openai/internal/test/checks"

"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"strconv"
"testing"
"time"
Expand All @@ -34,7 +35,7 @@ func TestFileUpload(t *testing.T) {
Purpose: "fine-tune",
}
_, err = client.CreateFile(ctx, req)
checks.NoError(t, err, "CreateFile erro")
checks.NoError(t, err, "CreateFile error")
}

// handleCreateFile Handles the images endpoint by the test server.
Expand Down Expand Up @@ -78,3 +79,50 @@ func handleCreateFile(w http.ResponseWriter, r *http.Request) {
resBytes, _ = json.Marshal(fileReq)
fmt.Fprint(w, string(resBytes))
}

func TestFileUploadWithFailingFormBuilder(t *testing.T) {
config := DefaultConfig("")
config.BaseURL = ""
client := NewClientWithConfig(config)
mockBuilder := &mockFormBuilder{}
client.createFormBuilder = func(io.Writer) formBuilder {
return mockBuilder
}

ctx := context.Background()
req := FileRequest{
FileName: "test.go",
FilePath: "client.go",
Purpose: "fine-tune",
}

mockError := fmt.Errorf("mockWriteField error")
mockBuilder.mockWriteField = func(string, string) error {
return mockError
}
_, err := client.CreateFile(ctx, req)
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")

mockError = fmt.Errorf("mockCreateFormFile error")
mockBuilder.mockWriteField = func(string, string) error {
return nil
}
mockBuilder.mockCreateFormFile = func(string, *os.File) error {
return mockError
}
_, err = client.CreateFile(ctx, req)
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")

mockError = fmt.Errorf("mockClose error")
mockBuilder.mockWriteField = func(string, string) error {
return nil
}
mockBuilder.mockCreateFormFile = func(string, *os.File) error {
return nil
}
mockBuilder.mockClose = func() error {
return mockError
}
_, err = client.CreateFile(ctx, req)
checks.ErrorIs(t, err, mockError, "CreateFile should return error if form builder fails")
}

0 comments on commit 33ca1da

Please sign in to comment.