File size: 2,584 Bytes
7def60a |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
package e2e_test
import (
"context"
"fmt"
"os"
"runtime"
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/ory/dockertest/v3"
"github.com/ory/dockertest/v3/docker"
"github.com/sashabaranov/go-openai"
)
var pool *dockertest.Pool
var resource *dockertest.Resource
var client *openai.Client
var containerImage = os.Getenv("LOCALAI_IMAGE")
var containerImageTag = os.Getenv("LOCALAI_IMAGE_TAG")
var modelsDir = os.Getenv("LOCALAI_MODELS_DIR")
var apiPort = os.Getenv("LOCALAI_API_PORT")
var apiEndpoint = os.Getenv("LOCALAI_API_ENDPOINT")
var apiKey = os.Getenv("LOCALAI_API_KEY")
func TestLocalAI(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "LocalAI E2E test suite")
}
var _ = BeforeSuite(func() {
if apiPort == "" {
apiPort = "8080"
}
var defaultConfig openai.ClientConfig
if apiEndpoint == "" {
startDockerImage()
defaultConfig = openai.DefaultConfig(apiKey)
apiEndpoint = "http://localhost:" + apiPort + "/v1" // So that other tests can reference this value safely.
defaultConfig.BaseURL = apiEndpoint
} else {
fmt.Println("Default ", apiEndpoint)
defaultConfig = openai.DefaultConfig(apiKey)
defaultConfig.BaseURL = apiEndpoint
}
// Wait for API to be ready
client = openai.NewClientWithConfig(defaultConfig)
Eventually(func() error {
_, err := client.ListModels(context.TODO())
return err
}, "20m").ShouldNot(HaveOccurred())
})
var _ = AfterSuite(func() {
if resource != nil {
Expect(pool.Purge(resource)).To(Succeed())
}
//dat, err := os.ReadFile(resource.Container.LogPath)
//Expect(err).To(Not(HaveOccurred()))
//Expect(string(dat)).To(ContainSubstring("GRPC Service Ready"))
//fmt.Println(string(dat))
})
var _ = AfterEach(func() {
//Expect(dbClient.Clear()).To(Succeed())
})
func startDockerImage() {
p, err := dockertest.NewPool("")
Expect(err).To(Not(HaveOccurred()))
Expect(p.Client.Ping()).To(Succeed())
pool = p
// get cwd
cwd, err := os.Getwd()
Expect(err).To(Not(HaveOccurred()))
md := cwd + "/models"
if modelsDir != "" {
md = modelsDir
}
proc := runtime.NumCPU()
options := &dockertest.RunOptions{
Repository: containerImage,
Tag: containerImageTag,
// Cmd: []string{"server", "/data"},
PortBindings: map[docker.Port][]docker.PortBinding{
"8080/tcp": []docker.PortBinding{{HostPort: apiPort}},
},
Env: []string{"MODELS_PATH=/models", "DEBUG=true", "THREADS=" + fmt.Sprint(proc)},
Mounts: []string{md + ":/models"},
}
r, err := pool.RunWithOptions(options)
Expect(err).To(Not(HaveOccurred()))
resource = r
}
|