From cc730268e7db0ca1cd281eac5452971b91effd06 Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Fri, 31 Jan 2025 11:11:44 -0500 Subject: [PATCH] Fix race condition waiting for status of ollama model pull success status --- .../org/springframework/ai/ollama/api/OllamaApiModelsIT.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaApiModelsIT.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaApiModelsIT.java index 6a0886289..17f5222df 100644 --- a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaApiModelsIT.java +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaApiModelsIT.java @@ -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));