Integration test improvements

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

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"})
} }

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,36 +174,51 @@ 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()
fp, err := os.Open(lifecycle.ServerLogFile) if t.Failed() {
if err != nil { fp, err := os.Open(lifecycle.ServerLogFile)
slog.Error("failed to open server log", "logfile", lifecycle.ServerLogFile, "error", err) if err != nil {
return slog.Error("failed to open server log", "logfile", lifecycle.ServerLogFile, "error", err)
return
}
data, err := io.ReadAll(fp)
if err != nil {
slog.Error("failed to read server log", "logfile", lifecycle.ServerLogFile, "error", err)
return
}
slog.Warn("SERVER LOG FOLLOWS")
os.Stderr.Write(data)
slog.Warn("END OF SERVER")
} }
data, err := io.ReadAll(fp) err = os.Remove(lifecycle.ServerLogFile)
if err != nil { if err != nil && !os.IsNotExist(err) {
slog.Error("failed to read server log", "logfile", lifecycle.ServerLogFile, "error", err) slog.Warn("failed to cleanup", "logfile", lifecycle.ServerLogFile, "error", err)
return
} }
slog.Warn("SERVER LOG FOLLOWS")
os.Stderr.Write(data)
slog.Warn("END OF SERVER")
}
err = os.Remove(lifecycle.ServerLogFile)
if err != nil && !os.IsNotExist(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))
} }