Skip to content

Commit 50db42f

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 c2103b0 commit 50db42f

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
@@ -157,6 +157,102 @@ public void vertexAiGeminiChatWithEmptyResponse() throws Exception {
157157
assertThat(this.retryListener.onErrorRetryCount).isEqualTo(1);
158158
}
159159

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

162258
int onErrorRetryCount = 0;

‎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)