Fix race condition waiting for status of ollama model pull success status

This commit is contained in:
Mark Pollack
2025-01-31 11:11:44 -05:00
parent a61437df6e
commit cc730268e7

View File

@@ -19,6 +19,7 @@ package org.springframework.ai.ollama.api;
import java.io.IOException;
import java.time.Duration;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -89,6 +90,10 @@ public class OllamaApiModelsIT extends BaseOllamaIT {
.block();
assertThat(progressResponses).isNotNull();
Awaitility.await().until(() -> {
OllamaApi.ProgressResponse progressResponse = progressResponses.get(progressResponses.size() - 1);
return progressResponse.status().equals("success");
});
assertThat(progressResponses.get(progressResponses.size() - 1))
.isEqualTo(new OllamaApi.ProgressResponse("success", null, null, null));