Integration test tuning (#12492)

Remove some flaky scenarios, and switch to chat for better reliability

integration/basic_test.go

@@ -17,16 +17,21 @@ func TestBlueSky(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 	// Set up the test data
-	req := api.GenerateRequest{
-		Model:  smol,
-		Prompt: blueSkyPrompt,
+	req := api.ChatRequest{
+		Model: smol,
+		Messages: []api.Message{
+			{
+				Role:    "user",
+				Content: blueSkyPrompt,
+			},
+		},
 		Stream: &stream,
 		Options: map[string]any{
 			"temperature": 0,
 			"seed":        123,
 		},
 	}
-	GenerateTestHelper(ctx, t, req, blueSkyExpected)
+	ChatTestHelper(ctx, t, req, blueSkyExpected)
 }
 
 func TestUnicode(t *testing.T) {
@@ -34,10 +39,15 @@ func TestUnicode(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
 	defer cancel()
 	// Set up the test data
-	req := api.GenerateRequest{
+	req := api.ChatRequest{
 		// DeepSeek has a Unicode tokenizer regex, making it a unicode torture test
-		Model:  "deepseek-coder-v2:16b-lite-instruct-q2_K", // TODO is there an ollama-engine model we can switch to and keep the coverage?
-		Prompt: "天空为什么是蓝色的?",                               // Why is the sky blue?
+		Model: "deepseek-coder-v2:16b-lite-instruct-q2_K", // TODO is there an ollama-engine model we can switch to and keep the coverage?
+		Messages: []api.Message{
+			{
+				Role:    "user",
+				Content: "天空为什么是蓝色的?", // Why is the sky blue?
+			},
+		},
 		Stream: &stream,
 		Options: map[string]any{
 			"temperature": 0,
@@ -57,9 +67,14 @@ func TestUnicode(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to load model %s: %s", req.Model, err)
 	}
+	defer func() {
+		// best effort unload once we're done with the model
+		client.Generate(ctx, &api.GenerateRequest{Model: req.Model, KeepAlive: &api.Duration{Duration: 0}}, func(rsp api.GenerateResponse) error { return nil })
+	}()
+
 	skipIfNotGPULoaded(ctx, t, client, req.Model, 100)
 
-	DoGenerate(ctx, t, client, req, []string{
+	DoChat(ctx, t, client, req, []string{
 		"散射", // scattering
 		"频率", // frequency
 	}, 120*time.Second, 120*time.Second)
@@ -69,9 +84,14 @@ func TestExtendedUnicodeOutput(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 	// Set up the test data
-	req := api.GenerateRequest{
-		Model:  "gemma2:2b",
-		Prompt: "Output some smily face emoji",
+	req := api.ChatRequest{
+		Model: "gemma2:2b",
+		Messages: []api.Message{
+			{
+				Role:    "user",
+				Content: "Output some smily face emoji",
+			},
+		},
 		Stream: &stream,
 		Options: map[string]any{
 			"temperature": 0,
@@ -83,7 +103,7 @@ func TestExtendedUnicodeOutput(t *testing.T) {
 	if err := PullIfMissing(ctx, client, req.Model); err != nil {
 		t.Fatal(err)
 	}
-	DoGenerate(ctx, t, client, req, []string{"😀", "😊", "😁", "😂", "😄", "😃"}, 120*time.Second, 120*time.Second)
+	DoChat(ctx, t, client, req, []string{"😀", "😊", "😁", "😂", "😄", "😃"}, 120*time.Second, 120*time.Second)
 }
 
 func TestUnicodeModelDir(t *testing.T) {
@@ -108,14 +128,19 @@ func TestUnicodeModelDir(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	req := api.GenerateRequest{
-		Model:  smol,
-		Prompt: blueSkyPrompt,
+	req := api.ChatRequest{
+		Model: smol,
+		Messages: []api.Message{
+			{
+				Role:    "user",
+				Content: blueSkyPrompt,
+			},
+		},
 		Stream: &stream,
 		Options: map[string]any{
 			"temperature": 0,
 			"seed":        123,
 		},
 	}
-	GenerateTestHelper(ctx, t, req, blueSkyExpected)
+	ChatTestHelper(ctx, t, req, blueSkyExpected)
 }