Skip to content

Commit b0ecd5e

Browse files
alxkmalxkm
committed
Add comprehensive test coverage for KeywordMetadataEnricher, ToolRuntimeHints, and Usage classes
Signed-off-by: Alex Klimenko <[email protected]>
1 parent e0ccc13 commit b0ecd5e

File tree

3 files changed

+158
-0
lines changed

3 files changed

+158
-0
lines changed

‎spring-ai-model/src/test/java/org/springframework/ai/aot/ToolRuntimeHintsTests.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,36 @@ void registerHintsWithNullClassLoader() {
4747
assertThatCode(() -> toolRuntimeHints.registerHints(runtimeHints, null)).doesNotThrowAnyException();
4848
}
4949

50+
@Test
51+
void registerHintsWithCustomClassLoader() {
52+
RuntimeHints runtimeHints = new RuntimeHints();
53+
ToolRuntimeHints toolRuntimeHints = new ToolRuntimeHints();
54+
ClassLoader customClassLoader = Thread.currentThread().getContextClassLoader();
55+
56+
toolRuntimeHints.registerHints(runtimeHints, customClassLoader);
57+
58+
assertThat(runtimeHints).matches(reflection().onType(DefaultToolCallResultConverter.class));
59+
}
60+
61+
@Test
62+
void registerHintsMultipleTimes() {
63+
RuntimeHints runtimeHints = new RuntimeHints();
64+
ToolRuntimeHints toolRuntimeHints = new ToolRuntimeHints();
65+
66+
toolRuntimeHints.registerHints(runtimeHints, null);
67+
toolRuntimeHints.registerHints(runtimeHints, null);
68+
69+
assertThat(runtimeHints).matches(reflection().onType(DefaultToolCallResultConverter.class));
70+
}
71+
72+
@Test
73+
void toolRuntimeHintsInstanceCreation() {
74+
assertThatCode(() -> new ToolRuntimeHints()).doesNotThrowAnyException();
75+
76+
ToolRuntimeHints hints1 = new ToolRuntimeHints();
77+
ToolRuntimeHints hints2 = new ToolRuntimeHints();
78+
79+
assertThat(hints1).isNotSameAs(hints2);
80+
}
81+
5082
}

‎spring-ai-model/src/test/java/org/springframework/ai/metadata/UsageTests.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,4 +87,28 @@ void totalTokensEqualsPromptTokensPlusGenerationTokens() {
8787
verifyUsage(usage);
8888
}
8989

90+
@Test
91+
void totalTokensHandlesZeroPromptTokens() {
92+
Usage usage = mockUsage(0, 1);
93+
94+
assertThat(usage.getTotalTokens()).isEqualTo(1);
95+
verifyUsage(usage);
96+
}
97+
98+
@Test
99+
void totalTokensHandlesZeroCompletionTokens() {
100+
Usage usage = mockUsage(1, 0);
101+
102+
assertThat(usage.getTotalTokens()).isEqualTo(1);
103+
verifyUsage(usage);
104+
}
105+
106+
@Test
107+
void totalTokensHandlesBothZeroTokens() {
108+
Usage usage = mockUsage(0, 0);
109+
110+
assertThat(usage.getTotalTokens()).isZero();
111+
verifyUsage(usage);
112+
}
113+
90114
}

‎spring-ai-model/src/test/java/org/springframework/ai/model/transformer/KeywordMetadataEnricherTest.java‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,4 +165,106 @@ private String getDefaultTemplatePromptText(int keywordCount, String documentCon
165165
return prompt.getContents();
166166
}
167167

168+
@Test
169+
void testApplyWithEmptyDocumentsList() {
170+
List<Document> emptyDocuments = List.of();
171+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 3);
172+
173+
keywordMetadataEnricher.apply(emptyDocuments);
174+
175+
verify(chatModel, never()).call(any(Prompt.class));
176+
}
177+
178+
@Test
179+
void testApplyWithSingleDocument() {
180+
List<Document> documents = List.of(new Document("single content"));
181+
given(chatModel.call(any(Prompt.class))).willReturn(new ChatResponse(
182+
List.of(new Generation(new AssistantMessage("single, keyword, test, document, content")))));
183+
184+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 5);
185+
keywordMetadataEnricher.apply(documents);
186+
187+
verify(chatModel, times(1)).call(promptCaptor.capture());
188+
assertThat(documents.get(0).getMetadata()).containsEntry(EXCERPT_KEYWORDS_METADATA_KEY,
189+
"single, keyword, test, document, content");
190+
}
191+
192+
@Test
193+
void testApplyWithDocumentContainingExistingMetadata() {
194+
Document document = new Document("content with existing metadata");
195+
document.getMetadata().put("existing_key", "existing_value");
196+
List<Document> documents = List.of(document);
197+
given(chatModel.call(any(Prompt.class)))
198+
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("new, keywords")))));
199+
200+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 2);
201+
keywordMetadataEnricher.apply(documents);
202+
203+
assertThat(documents.get(0).getMetadata()).containsEntry("existing_key", "existing_value");
204+
assertThat(documents.get(0).getMetadata()).containsEntry(EXCERPT_KEYWORDS_METADATA_KEY, "new, keywords");
205+
}
206+
207+
@Test
208+
void testApplyWithEmptyStringResponse() {
209+
List<Document> documents = List.of(new Document("content"));
210+
given(chatModel.call(any(Prompt.class)))
211+
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("")))));
212+
213+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 3);
214+
keywordMetadataEnricher.apply(documents);
215+
216+
assertThat(documents.get(0).getMetadata()).containsEntry(EXCERPT_KEYWORDS_METADATA_KEY, "");
217+
}
218+
219+
@Test
220+
void testApplyWithWhitespaceOnlyResponse() {
221+
List<Document> documents = List.of(new Document("content"));
222+
given(chatModel.call(any(Prompt.class)))
223+
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage(" \n\t ")))));
224+
225+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 3);
226+
keywordMetadataEnricher.apply(documents);
227+
228+
assertThat(documents.get(0).getMetadata()).containsEntry(EXCERPT_KEYWORDS_METADATA_KEY, " \n\t ");
229+
}
230+
231+
@Test
232+
void testApplyOverwritesExistingKeywords() {
233+
Document document = new Document("content");
234+
document.getMetadata().put(EXCERPT_KEYWORDS_METADATA_KEY, "old, keywords");
235+
List<Document> documents = List.of(document);
236+
given(chatModel.call(any(Prompt.class)))
237+
.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("new, keywords")))));
238+
239+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 2);
240+
keywordMetadataEnricher.apply(documents);
241+
242+
assertThat(documents.get(0).getMetadata()).containsEntry(EXCERPT_KEYWORDS_METADATA_KEY, "new, keywords");
243+
}
244+
245+
@Test
246+
void testBuilderWithBothKeywordCountAndTemplate() {
247+
PromptTemplate customTemplate = new PromptTemplate(CUSTOM_TEMPLATE);
248+
249+
KeywordMetadataEnricher enricher = builder(chatModel).keywordCount(5).keywordsTemplate(customTemplate).build();
250+
251+
assertThat(enricher.getKeywordsTemplate()).isEqualTo(customTemplate);
252+
}
253+
254+
@Test
255+
void testApplyWithSpecialCharactersInContent() {
256+
List<Document> documents = List.of(new Document("Content with special chars: @#$%^&*()"));
257+
given(chatModel.call(any(Prompt.class))).willReturn(
258+
new ChatResponse(List.of(new Generation(new AssistantMessage("special, characters, content")))));
259+
260+
KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 3);
261+
keywordMetadataEnricher.apply(documents);
262+
263+
verify(chatModel, times(1)).call(promptCaptor.capture());
264+
assertThat(promptCaptor.getValue().getUserMessage().getText())
265+
.contains("Content with special chars: @#$%^&*()");
266+
assertThat(documents.get(0).getMetadata()).containsEntry(EXCERPT_KEYWORDS_METADATA_KEY,
267+
"special, characters, content");
268+
}
269+
168270
}

0 commit comments

Comments
 (0)