baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Integration test improvements

Browse source 
Cleaner shutdown logic, a bit of response hardening
This commit is contained in:
Daniel Hiltgen 2024-03-27 11:02:06 -07:00
parent 0a0e9f3e0f
commit 4fec5816d6
2 changed files with 34 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
integration/basic_test.go
View file

@ -24,5 +24,5 @@ func TestOrcaMiniBlueSky(t *testing.T) {
"seed": 123, "seed": 123,
}, },
} }
GenerateTestHelper(ctx, t, &http.Client{}, req, []string{"rayleigh"}) GenerateTestHelper(ctx, t, &http.Client{}, req, []string{"rayleigh", "scattering"})
} }

21
integration/utils_test.go
View file

@ -126,7 +126,7 @@ func StartServer(ctx context.Context, ollamaHost string) error {
} }
func PullIfMissing(ctx context.Context, client *http.Client, scheme, testEndpoint, modelName string) error { func PullIfMissing(ctx context.Context, client *http.Client, scheme, testEndpoint, modelName string) error {
slog.Debug("checking status of model", "model", modelName) slog.Info("checking status of model", "model", modelName)
showReq := &api.ShowRequest{Name: modelName} showReq := &api.ShowRequest{Name: modelName}
requestJSON, err := json.Marshal(showReq) requestJSON, err := json.Marshal(showReq)
if err != nil { if err != nil {
@ -174,14 +174,21 @@ func PullIfMissing(ctx context.Context, client *http.Client, scheme, testEndpoin
return nil return nil
} }
var serverProcMutex sync.Mutex
func GenerateTestHelper(ctx context.Context, t *testing.T, client *http.Client, genReq api.GenerateRequest, anyResp []string) { func GenerateTestHelper(ctx context.Context, t *testing.T, client *http.Client, genReq api.GenerateRequest, anyResp []string) {
// TODO maybe stuff in an init routine?
lifecycle.InitLogging()
requestJSON, err := json.Marshal(genReq) requestJSON, err := json.Marshal(genReq)
if err != nil { if err != nil {
t.Fatalf("Error serializing request: %v", err) t.Fatalf("Error serializing request: %v", err)
} }
defer func() { defer func() {
if t.Failed() && os.Getenv("OLLAMA_TEST_EXISTING") == "" { if os.Getenv("OLLAMA_TEST_EXISTING") == "" {
// TODO defer serverProcMutex.Unlock()
if t.Failed() {
fp, err := os.Open(lifecycle.ServerLogFile) fp, err := os.Open(lifecycle.ServerLogFile)
if err != nil { if err != nil {
slog.Error("failed to open server log", "logfile", lifecycle.ServerLogFile, "error", err) slog.Error("failed to open server log", "logfile", lifecycle.ServerLogFile, "error", err)
@ -200,10 +207,18 @@ func GenerateTestHelper(ctx context.Context, t *testing.T, client *http.Client,
if err != nil && !os.IsNotExist(err) { if err != nil && !os.IsNotExist(err) {
slog.Warn("failed to cleanup", "logfile", lifecycle.ServerLogFile, "error", err) slog.Warn("failed to cleanup", "logfile", lifecycle.ServerLogFile, "error", err)
} }
}
}() }()
scheme, testEndpoint := GetTestEndpoint() scheme, testEndpoint := GetTestEndpoint()
if os.Getenv("OLLAMA_TEST_EXISTING") == "" { if os.Getenv("OLLAMA_TEST_EXISTING") == "" {
serverProcMutex.Lock()
fp, err := os.CreateTemp("", "ollama-server-*.log")
if err != nil {
t.Fatalf("failed to generate log file: %s", err)
}
lifecycle.ServerLogFile = fp.Name()
fp.Close()
assert.NoError(t, StartServer(ctx, testEndpoint)) assert.NoError(t, StartServer(ctx, testEndpoint))
} }