models_test.go

package openai_test

import (
"context"
"encoding/json"
"fmt"
"net/http"
"os"
"testing"
"time"

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

consttestFineTuneModelID="fine-tune-model-id"

// TestListModels Tests the list models endpoint of the API using the mocked server.
funcTestListModels(t*testing.T) {
client, server, teardown:=setupOpenAITestServer()
deferteardown()
server.RegisterHandler("/v1/models", handleListModelsEndpoint)
_, err:=client.ListModels(context.Background())
checks.NoError(t, err, "ListModels error")
}

funcTestAzureListModels(t*testing.T) {
client, server, teardown:=setupAzureTestServer()
deferteardown()
server.RegisterHandler("/openai/models", handleListModelsEndpoint)
_, err:=client.ListModels(context.Background())
checks.NoError(t, err, "ListModels error")
}

// handleListModelsEndpoint Handles the list models endpoint by the test server.
funchandleListModelsEndpoint(w http.ResponseWriter, _*http.Request) {
resBytes, _:=json.Marshal(openai.ModelsList{})
fmt.Fprintln(w, string(resBytes))
}

// TestGetModel Tests the retrieve model endpoint of the API using the mocked server.
funcTestGetModel(t*testing.T) {
client, server, teardown:=setupOpenAITestServer()
deferteardown()
server.RegisterHandler("/v1/models/text-davinci-003", handleGetModelEndpoint)
_, err:=client.GetModel(context.Background(), "text-davinci-003")
checks.NoError(t, err, "GetModel error")
}

funcTestAzureGetModel(t*testing.T) {
client, server, teardown:=setupAzureTestServer()
deferteardown()
server.RegisterHandler("/openai/models/text-davinci-003", handleGetModelEndpoint)
_, err:=client.GetModel(context.Background(), "text-davinci-003")
checks.NoError(t, err, "GetModel error")
}

// handleGetModelsEndpoint Handles the get model endpoint by the test server.
funchandleGetModelEndpoint(w http.ResponseWriter, _*http.Request) {
resBytes, _:=json.Marshal(openai.Model{})
fmt.Fprintln(w, string(resBytes))
}

funcTestGetModelReturnTimeoutError(t*testing.T) {
client, server, teardown:=setupOpenAITestServer()
deferteardown()
server.RegisterHandler("/v1/models/text-davinci-003", func(http.ResponseWriter, *http.Request) {
time.Sleep(10*time.Nanosecond)
 })
ctx:=context.Background()
ctx, cancel:=context.WithTimeout(ctx, time.Nanosecond)
defercancel()

_, err:=client.GetModel(ctx, "text-davinci-003")
iferr==nil {
t.Fatal("Did not return error")
 }
if!os.IsTimeout(err) {
t.Fatal("Did not return timeout error")
 }
}

funcTestDeleteFineTuneModel(t*testing.T) {
client, server, teardown:=setupOpenAITestServer()
deferteardown()
server.RegisterHandler("/v1/models/"+testFineTuneModelID, handleDeleteFineTuneModelEndpoint)
_, err:=client.DeleteFineTuneModel(context.Background(), testFineTuneModelID)
checks.NoError(t, err, "DeleteFineTuneModel error")
}

funchandleDeleteFineTuneModelEndpoint(w http.ResponseWriter, _*http.Request) {
resBytes, _:=json.Marshal(openai.FineTuneModelDeleteResponse{})
fmt.Fprintln(w, string(resBytes))
}