Skip to content

Commit 30fb517

Browse files
alxkmalxkm
committed
test: Add comprehensive test coverage for ModelObservationContext and ModelUsageMetricsGenerator
Signed-off-by: Alex Klimenko <[email protected]>
1 parent e0ccc13 commit 30fb517

File tree

2 files changed

+187
-0
lines changed

2 files changed

+187
-0
lines changed

‎spring-ai-model/src/test/java/org/springframework/ai/model/observation/ModelObservationContextTests.java‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,4 +100,157 @@ void whenResponseIsNullThenThrow() {
100100
.hasMessageContaining("response cannot be null");
101101
}
102102

103+
@Test
104+
void whenEmptyOperationTypeThenThrow() {
105+
assertThatThrownBy(() -> new ModelObservationContext<String, String>("test request",
106+
AiOperationMetadata.builder().operationType("").provider(AiProvider.OLLAMA.value()).build()))
107+
.isInstanceOf(IllegalArgumentException.class);
108+
}
109+
110+
@Test
111+
void whenEmptyProviderThenThrow() {
112+
assertThatThrownBy(() -> new ModelObservationContext<String, String>("test request",
113+
AiOperationMetadata.builder().operationType(AiOperationType.CHAT.value()).provider("").build()))
114+
.isInstanceOf(IllegalArgumentException.class);
115+
}
116+
117+
@Test
118+
void whenDifferentProvidersThenReturn() {
119+
var ollamaContext = new ModelObservationContext<String, String>("test request",
120+
AiOperationMetadata.builder()
121+
.operationType(AiOperationType.CHAT.value())
122+
.provider(AiProvider.OLLAMA.value())
123+
.build());
124+
125+
var openaiContext = new ModelObservationContext<String, String>("test request",
126+
AiOperationMetadata.builder()
127+
.operationType(AiOperationType.CHAT.value())
128+
.provider(AiProvider.OPENAI.value())
129+
.build());
130+
131+
var anthropicContext = new ModelObservationContext<String, String>("test request",
132+
AiOperationMetadata.builder()
133+
.operationType(AiOperationType.CHAT.value())
134+
.provider(AiProvider.ANTHROPIC.value())
135+
.build());
136+
137+
assertThat(ollamaContext).isNotNull();
138+
assertThat(openaiContext).isNotNull();
139+
assertThat(anthropicContext).isNotNull();
140+
}
141+
142+
@Test
143+
void whenComplexObjectTypesAreUsedThenReturn() {
144+
var observationContext = new ModelObservationContext<Integer, Boolean>(12345,
145+
AiOperationMetadata.builder()
146+
.operationType(AiOperationType.CHAT.value())
147+
.provider(AiProvider.OLLAMA.value())
148+
.build());
149+
observationContext.setResponse(true);
150+
151+
assertThat(observationContext).isNotNull();
152+
}
153+
154+
@Test
155+
void whenGetRequestThenReturn() {
156+
var testRequest = "test request content";
157+
var observationContext = new ModelObservationContext<String, String>(testRequest,
158+
AiOperationMetadata.builder()
159+
.operationType(AiOperationType.CHAT.value())
160+
.provider(AiProvider.OLLAMA.value())
161+
.build());
162+
163+
assertThat(observationContext.getRequest()).isEqualTo(testRequest);
164+
}
165+
166+
@Test
167+
void whenGetResponseBeforeSettingThenReturnNull() {
168+
var observationContext = new ModelObservationContext<String, String>("test request",
169+
AiOperationMetadata.builder()
170+
.operationType(AiOperationType.CHAT.value())
171+
.provider(AiProvider.OLLAMA.value())
172+
.build());
173+
174+
assertThat(observationContext.getResponse()).isNull();
175+
}
176+
177+
@Test
178+
void whenGetResponseAfterSettingThenReturn() {
179+
var testResponse = "test response content";
180+
var observationContext = new ModelObservationContext<String, String>("test request",
181+
AiOperationMetadata.builder()
182+
.operationType(AiOperationType.CHAT.value())
183+
.provider(AiProvider.OLLAMA.value())
184+
.build());
185+
observationContext.setResponse(testResponse);
186+
187+
assertThat(observationContext.getResponse()).isEqualTo(testResponse);
188+
}
189+
190+
@Test
191+
void whenGetOperationMetadataThenReturn() {
192+
var metadata = AiOperationMetadata.builder()
193+
.operationType(AiOperationType.EMBEDDING.value())
194+
.provider(AiProvider.OPENAI.value())
195+
.build();
196+
var observationContext = new ModelObservationContext<String, String>("test request", metadata);
197+
198+
assertThat(observationContext.getOperationMetadata()).isEqualTo(metadata);
199+
}
200+
201+
@Test
202+
void whenSetResponseMultipleTimesThenLastValueWins() {
203+
var observationContext = new ModelObservationContext<String, String>("test request",
204+
AiOperationMetadata.builder()
205+
.operationType(AiOperationType.CHAT.value())
206+
.provider(AiProvider.OLLAMA.value())
207+
.build());
208+
209+
observationContext.setResponse("first response");
210+
observationContext.setResponse("second response");
211+
observationContext.setResponse("final response");
212+
213+
assertThat(observationContext.getResponse()).isEqualTo("final response");
214+
}
215+
216+
@Test
217+
void whenWhitespaceOnlyOperationTypeThenThrow() {
218+
assertThatThrownBy(() -> new ModelObservationContext<String, String>("test request",
219+
AiOperationMetadata.builder().operationType(" ").provider(AiProvider.OLLAMA.value()).build()))
220+
.isInstanceOf(IllegalArgumentException.class)
221+
.hasMessageContaining("operationType cannot be null or empty");
222+
}
223+
224+
@Test
225+
void whenWhitespaceOnlyProviderThenThrow() {
226+
assertThatThrownBy(() -> new ModelObservationContext<String, String>("test request",
227+
AiOperationMetadata.builder().operationType(AiOperationType.CHAT.value()).provider(" ").build()))
228+
.isInstanceOf(IllegalArgumentException.class)
229+
.hasMessageContaining("provider cannot be null or empty");
230+
}
231+
232+
@Test
233+
void whenEmptyStringRequestThenReturn() {
234+
var observationContext = new ModelObservationContext<String, String>("",
235+
AiOperationMetadata.builder()
236+
.operationType(AiOperationType.CHAT.value())
237+
.provider(AiProvider.OLLAMA.value())
238+
.build());
239+
240+
assertThat(observationContext).isNotNull();
241+
assertThat(observationContext.getRequest()).isEqualTo("");
242+
}
243+
244+
@Test
245+
void whenEmptyStringResponseThenReturn() {
246+
var observationContext = new ModelObservationContext<String, String>("test request",
247+
AiOperationMetadata.builder()
248+
.operationType(AiOperationType.CHAT.value())
249+
.provider(AiProvider.OLLAMA.value())
250+
.build());
251+
observationContext.setResponse("");
252+
253+
assertThat(observationContext.getResponse()).isEqualTo("");
254+
}
255+
103256
}

‎spring-ai-model/src/test/java/org/springframework/ai/model/observation/ModelUsageMetricsGeneratorTests.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,4 +123,38 @@ public Map<String, Integer> getNativeUsage() {
123123

124124
}
125125

126+
@Test
127+
void whenZeroTokenUsageThenMetrics() {
128+
var meterRegistry = new SimpleMeterRegistry();
129+
var usage = new TestUsage(0, 0, 0);
130+
ModelUsageMetricsGenerator.generate(usage, buildContext(), meterRegistry);
131+
132+
assertThat(meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value()).meters()).hasSize(3);
133+
assertThat(meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
134+
.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.INPUT.value())
135+
.counter()
136+
.count()).isEqualTo(0);
137+
assertThat(meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
138+
.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.OUTPUT.value())
139+
.counter()
140+
.count()).isEqualTo(0);
141+
assertThat(meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
142+
.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.TOTAL.value())
143+
.counter()
144+
.count()).isEqualTo(0);
145+
}
146+
147+
@Test
148+
void whenBothPromptAndGenerationNullThenOnlyTotalMetric() {
149+
var meterRegistry = new SimpleMeterRegistry();
150+
var usage = new TestUsage(null, null, 100);
151+
ModelUsageMetricsGenerator.generate(usage, buildContext(), meterRegistry);
152+
153+
assertThat(meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value()).meters()).hasSize(1);
154+
assertThat(meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
155+
.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.TOTAL.value())
156+
.counter()
157+
.count()).isEqualTo(100);
158+
}
159+
126160
}

0 commit comments

Comments
 (0)