Skip to content

Commit 67bc806

Browse files
alxkmalxkm
committed
test: Add comprehensive test coverage for PromptMetadata and VertexAI embedding retry tests
Co-authored-by: Oleksandr Klymenko <[email protected]> Signed-off-by: Oleksandr Klymenko <[email protected]>
1 parent ad93613 commit 67bc806

File tree

3 files changed

+178
-0
lines changed

3 files changed

+178
-0
lines changed

‎models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingRetryTests.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,19 @@ public void vertexAiEmbeddingNonTransientError() {
143143
verify(this.mockPredictionServiceClient, times(1)).predict(any());
144144
}
145145

146+
@Test
147+
public void vertexAiEmbeddingWithEmptyTextList() {
148+
PredictResponse emptyResponse = PredictResponse.newBuilder().build();
149+
given(this.mockPredictionServiceClient.predict(any())).willReturn(emptyResponse);
150+
151+
EmbeddingOptions options = VertexAiTextEmbeddingOptions.builder().model("model").build();
152+
EmbeddingResponse result = this.embeddingModel.call(new EmbeddingRequest(List.of(), options));
153+
154+
assertThat(result).isNotNull();
155+
// Behavior depends on implementation - might be empty results or exception
156+
verify(this.mockPredictionServiceClient, times(1)).predict(any());
157+
}
158+
146159
private static class TestRetryListener implements RetryListener {
147160

148161
int onErrorRetryCount = 0;

‎models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiRetryTests.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,4 +176,100 @@ public void vertexAiGeminiChatWithEmptyResponse() throws Exception {
176176
assertThat(this.retryListener.onErrorRetryCount).isEqualTo(1);
177177
}
178178

179+
@Test
180+
public void vertexAiGeminiChatMaxRetriesExceeded() throws Exception {
181+
// Test that after max retries, the exception is propagated
182+
given(this.mockGenerativeModel.generateContent(any(List.class)))
183+
.willThrow(new TransientAiException("Persistent Error"))
184+
.willThrow(new TransientAiException("Persistent Error"))
185+
.willThrow(new TransientAiException("Persistent Error"))
186+
.willThrow(new TransientAiException("Persistent Error"));
187+
188+
// Should throw the last TransientAiException after exhausting retries
189+
assertThrows(TransientAiException.class, () -> this.chatModel.call(new Prompt("test prompt")));
190+
191+
// Verify retry attempts were made
192+
assertThat(this.retryListener.onErrorRetryCount).isGreaterThan(0);
193+
}
194+
195+
@Test
196+
public void vertexAiGeminiChatWithMultipleCandidatesResponse() throws Exception {
197+
// Test response with multiple candidates
198+
GenerateContentResponse multiCandidateResponse = GenerateContentResponse.newBuilder()
199+
.addCandidates(Candidate.newBuilder()
200+
.setContent(Content.newBuilder().addParts(Part.newBuilder().setText("First candidate").build()).build())
201+
.build())
202+
.addCandidates(Candidate.newBuilder()
203+
.setContent(
204+
Content.newBuilder().addParts(Part.newBuilder().setText("Second candidate").build()).build())
205+
.build())
206+
.build();
207+
208+
given(this.mockGenerativeModel.generateContent(any(List.class)))
209+
.willThrow(new TransientAiException("Temporary failure"))
210+
.willReturn(multiCandidateResponse);
211+
212+
ChatResponse result = this.chatModel.call(new Prompt("test prompt"));
213+
214+
assertThat(result).isNotNull();
215+
// Assuming the implementation uses the first candidate
216+
assertThat(result.getResult().getOutput().getText()).isEqualTo("First candidate");
217+
assertThat(this.retryListener.onSuccessRetryCount).isEqualTo(1);
218+
}
219+
220+
@Test
221+
public void vertexAiGeminiChatWithNullPrompt() throws Exception {
222+
// Test handling of null prompt
223+
Prompt prompt = null;
224+
assertThrows(Exception.class, () -> this.chatModel.call(prompt));
225+
226+
// Should not trigger any retries for validation errors
227+
assertThat(this.retryListener.onErrorRetryCount).isEqualTo(0);
228+
assertThat(this.retryListener.onSuccessRetryCount).isEqualTo(0);
229+
}
230+
231+
@Test
232+
public void vertexAiGeminiChatWithEmptyPrompt() throws Exception {
233+
// Test handling of empty prompt
234+
GenerateContentResponse mockedResponse = GenerateContentResponse.newBuilder()
235+
.addCandidates(Candidate.newBuilder()
236+
.setContent(Content.newBuilder()
237+
.addParts(Part.newBuilder().setText("Response to empty prompt").build())
238+
.build())
239+
.build())
240+
.build();
241+
242+
given(this.mockGenerativeModel.generateContent(any(List.class))).willReturn(mockedResponse);
243+
244+
ChatResponse result = this.chatModel.call(new Prompt(""));
245+
246+
assertThat(result).isNotNull();
247+
assertThat(result.getResult().getOutput().getText()).isEqualTo("Response to empty prompt");
248+
assertThat(this.retryListener.onSuccessRetryCount).isEqualTo(0);
249+
}
250+
251+
@Test
252+
public void vertexAiGeminiChatAlternatingErrorsAndSuccess() throws Exception {
253+
// Test pattern of error -> success -> error -> success
254+
GenerateContentResponse successResponse = GenerateContentResponse.newBuilder()
255+
.addCandidates(Candidate.newBuilder()
256+
.setContent(Content.newBuilder()
257+
.addParts(Part.newBuilder().setText("Success after alternating errors").build())
258+
.build())
259+
.build())
260+
.build();
261+
262+
given(this.mockGenerativeModel.generateContent(any(List.class)))
263+
.willThrow(new TransientAiException("First error"))
264+
.willThrow(new TransientAiException("Second error"))
265+
.willReturn(successResponse);
266+
267+
ChatResponse result = this.chatModel.call(new Prompt("test prompt"));
268+
269+
assertThat(result).isNotNull();
270+
assertThat(result.getResult().getOutput().getText()).isEqualTo("Success after alternating errors");
271+
assertThat(this.retryListener.onSuccessRetryCount).isEqualTo(2);
272+
assertThat(this.retryListener.onErrorRetryCount).isEqualTo(2);
273+
}
274+
179275
}

‎spring-ai-client-chat/src/test/java/org/springframework/ai/metadata/PromptMetadataTests.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,4 +106,73 @@ void fromPromptIndexAndContentFilterMetadata() {
106106
assertThat(promptFilterMetadata.<String>getContentFilterMetadata()).isEqualTo("{ content-sentiment: 'SAFE' }");
107107
}
108108

109+
@Test
110+
void promptMetadataWithEmptyFiltersArray() {
111+
PromptMetadata promptMetadata = PromptMetadata.of();
112+
113+
assertThat(promptMetadata).isNotNull();
114+
assertThat(promptMetadata).isEmpty();
115+
}
116+
117+
@Test
118+
void promptMetadataWithMultipleFilters() {
119+
PromptFilterMetadata filter1 = mockPromptFilterMetadata(0);
120+
PromptFilterMetadata filter2 = mockPromptFilterMetadata(1);
121+
PromptFilterMetadata filter3 = mockPromptFilterMetadata(2);
122+
PromptFilterMetadata filter4 = mockPromptFilterMetadata(3);
123+
124+
PromptMetadata promptMetadata = PromptMetadata.of(filter1, filter2, filter3, filter4);
125+
126+
assertThat(promptMetadata).isNotNull();
127+
assertThat(promptMetadata).hasSize(4);
128+
assertThat(promptMetadata).containsExactly(filter1, filter2, filter3, filter4);
129+
}
130+
131+
@Test
132+
void promptMetadataWithDuplicateIndices() {
133+
PromptFilterMetadata filter1 = mockPromptFilterMetadata(1);
134+
PromptFilterMetadata filter2 = mockPromptFilterMetadata(1);
135+
136+
PromptMetadata promptMetadata = PromptMetadata.of(filter1, filter2);
137+
138+
assertThat(promptMetadata).isNotNull();
139+
assertThat(promptMetadata).hasSize(2);
140+
141+
assertThat(promptMetadata.findByPromptIndex(1).orElse(null)).isEqualTo(filter1);
142+
}
143+
144+
@Test
145+
void promptFilterMetadataWithEmptyContentFilter() {
146+
PromptFilterMetadata promptFilterMetadata = PromptFilterMetadata.from(0, "");
147+
148+
assertThat(promptFilterMetadata).isNotNull();
149+
assertThat(promptFilterMetadata.getPromptIndex()).isZero();
150+
assertThat(promptFilterMetadata.<String>getContentFilterMetadata()).isEmpty();
151+
}
152+
153+
@Test
154+
void promptMetadataSize() {
155+
PromptFilterMetadata filter1 = mockPromptFilterMetadata(0);
156+
PromptFilterMetadata filter2 = mockPromptFilterMetadata(1);
157+
158+
PromptMetadata empty = PromptMetadata.empty();
159+
PromptMetadata single = PromptMetadata.of(filter1);
160+
PromptMetadata multiple = PromptMetadata.of(filter1, filter2);
161+
162+
assertThat(empty).hasSize(0);
163+
assertThat(single).hasSize(1);
164+
assertThat(multiple).hasSize(2);
165+
}
166+
167+
@Test
168+
void promptMetadataImmutability() {
169+
PromptFilterMetadata filter1 = mockPromptFilterMetadata(0);
170+
PromptFilterMetadata filter2 = mockPromptFilterMetadata(1);
171+
172+
PromptMetadata promptMetadata = PromptMetadata.of(filter1, filter2);
173+
174+
assertThat(promptMetadata).isNotNull();
175+
assertThat(promptMetadata).hasSize(2);
176+
}
177+
109178
}

0 commit comments

Comments
 (0)