| 
18 | 18 | 
 
  | 
19 | 19 | import io.micrometer.tracing.Tracer;  | 
20 | 20 | import org.junit.jupiter.api.Test;  | 
 | 21 | +import org.junit.jupiter.api.extension.ExtendWith;  | 
 | 22 | +import org.springframework.ai.chat.client.observation.ChatClientObservationContext;  | 
21 | 23 | import org.springframework.ai.chat.client.observation.ChatClientPromptContentObservationHandler;  | 
 | 24 | +import org.springframework.ai.observation.TracingAwareLoggingObservationHandler;  | 
22 | 25 | import org.springframework.boot.autoconfigure.AutoConfigurations;  | 
23 | 26 | import org.springframework.boot.test.context.FilteredClassLoader;  | 
24 | 27 | import org.springframework.boot.test.context.runner.ApplicationContextRunner;  | 
 | 28 | +import org.springframework.boot.test.system.CapturedOutput;  | 
 | 29 | +import org.springframework.boot.test.system.OutputCaptureExtension;  | 
 | 30 | +import org.springframework.context.annotation.Bean;  | 
 | 31 | +import org.springframework.context.annotation.Configuration;  | 
25 | 32 | 
 
  | 
26 | 33 | import static org.assertj.core.api.Assertions.assertThat;  | 
 | 34 | +import static org.mockito.Mockito.mock;  | 
27 | 35 | 
 
  | 
28 | 36 | /**  | 
29 | 37 |  * Unit tests for {@link ChatClientAutoConfiguration} observability support.  | 
 | 
32 | 40 |  * @author Thomas Vitale  | 
33 | 41 |  * @author Jonatan Ivanov  | 
34 | 42 |  */  | 
 | 43 | +@ExtendWith(OutputCaptureExtension.class)  | 
35 | 44 | class ChatClientObservationAutoConfigurationTests {  | 
36 | 45 | 
 
  | 
37 | 46 | 	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()  | 
38 | 47 | 		.withConfiguration(AutoConfigurations.of(ChatClientAutoConfiguration.class));  | 
39 | 48 | 
 
  | 
40 | 49 | 	@Test  | 
41 |  | -	void promptContentHandlerDefault() {  | 
42 |  | -		this.contextRunner  | 
43 |  | -			.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class));  | 
 | 50 | +	void promptContentHandlerNoTracer() {  | 
 | 51 | +		this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))  | 
 | 52 | +			.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)  | 
 | 53 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
 | 54 | +	}  | 
 | 55 | + | 
 | 56 | +	@Test  | 
 | 57 | +	void promptContentHandlerWithTracer() {  | 
 | 58 | +		this.contextRunner.withUserConfiguration(TracerConfiguration.class)  | 
 | 59 | +			.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)  | 
 | 60 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
44 | 61 | 	}  | 
45 | 62 | 
 
  | 
46 | 63 | 	@Test  | 
47 |  | -	void promptContentHandlerEnabled() {  | 
 | 64 | +	void promptContentHandlerEnabledNoTracer(CapturedOutput output) {  | 
48 | 65 | 		this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))  | 
49 | 66 | 			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")  | 
50 |  | -			.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class));  | 
 | 67 | +			.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class)  | 
 | 68 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
 | 69 | +		assertThat(output).contains(  | 
 | 70 | +				"You have enabled logging out the ChatClient prompt content with the risk of exposing sensitive or private information. Please, be careful!");  | 
 | 71 | +	}  | 
 | 72 | + | 
 | 73 | +	@Test  | 
 | 74 | +	void promptContentHandlerEnabledWithTracer(CapturedOutput output) {  | 
 | 75 | +		this.contextRunner.withUserConfiguration(TracerConfiguration.class)  | 
 | 76 | +			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")  | 
 | 77 | +			.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)  | 
 | 78 | +				.hasSingleBean(TracingAwareLoggingObservationHandler.class));  | 
 | 79 | +		assertThat(output).contains(  | 
 | 80 | +				"You have enabled logging out the ChatClient prompt content with the risk of exposing sensitive or private information. Please, be careful!");  | 
 | 81 | +	}  | 
 | 82 | + | 
 | 83 | +	@Test  | 
 | 84 | +	void promptContentHandlerDisabledNoTracer() {  | 
 | 85 | +		this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))  | 
 | 86 | +			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=false")  | 
 | 87 | +			.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)  | 
 | 88 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
 | 89 | +	}  | 
 | 90 | + | 
 | 91 | +	@Test  | 
 | 92 | +	void promptContentHandlerDisabledWithTracer() {  | 
 | 93 | +		this.contextRunner.withUserConfiguration(TracerConfiguration.class)  | 
 | 94 | +			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=false")  | 
 | 95 | +			.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)  | 
 | 96 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
 | 97 | +	}  | 
 | 98 | + | 
 | 99 | +	@Test  | 
 | 100 | +	void customChatClientPromptContentObservationHandlerNoTracer() {  | 
 | 101 | +		this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))  | 
 | 102 | +			.withUserConfiguration(CustomChatClientPromptContentObservationHandlerConfiguration.class)  | 
 | 103 | +			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")  | 
 | 104 | +			.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class)  | 
 | 105 | +				.hasBean("customChatClientPromptContentObservationHandler")  | 
 | 106 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
 | 107 | +	}  | 
 | 108 | + | 
 | 109 | +	@Test  | 
 | 110 | +	void customChatClientPromptContentObservationHandlerWithTracer() {  | 
 | 111 | +		this.contextRunner.withUserConfiguration(TracerConfiguration.class)  | 
 | 112 | +			.withUserConfiguration(CustomChatClientPromptContentObservationHandlerConfiguration.class)  | 
 | 113 | +			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")  | 
 | 114 | +			.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class)  | 
 | 115 | +				.hasBean("customChatClientPromptContentObservationHandler")  | 
 | 116 | +				.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));  | 
 | 117 | +	}  | 
 | 118 | + | 
 | 119 | +	@Test  | 
 | 120 | +	void customTracingAwareLoggingObservationHandler() {  | 
 | 121 | +		this.contextRunner.withUserConfiguration(TracerConfiguration.class)  | 
 | 122 | +			.withUserConfiguration(CustomTracingAwareLoggingObservationHandlerConfiguration.class)  | 
 | 123 | +			.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")  | 
 | 124 | +			.run(context -> {  | 
 | 125 | +				assertThat(context).hasSingleBean(TracingAwareLoggingObservationHandler.class)  | 
 | 126 | +					.hasBean("chatClientPromptContentObservationHandler")  | 
 | 127 | +					.doesNotHaveBean(ChatClientPromptContentObservationHandler.class);  | 
 | 128 | +				assertThat(context.getBean(TracingAwareLoggingObservationHandler.class))  | 
 | 129 | +					.isSameAs(CustomTracingAwareLoggingObservationHandlerConfiguration.handlerInstance);  | 
 | 130 | +			});  | 
 | 131 | +	}  | 
 | 132 | + | 
 | 133 | +	@Configuration(proxyBeanMethods = false)  | 
 | 134 | +	static class TracerConfiguration {  | 
 | 135 | + | 
 | 136 | +		@Bean  | 
 | 137 | +		Tracer tracer() {  | 
 | 138 | +			return mock(Tracer.class);  | 
 | 139 | +		}  | 
 | 140 | + | 
 | 141 | +	}  | 
 | 142 | + | 
 | 143 | +	@Configuration(proxyBeanMethods = false)  | 
 | 144 | +	static class CustomChatClientPromptContentObservationHandlerConfiguration {  | 
 | 145 | + | 
 | 146 | +		@Bean  | 
 | 147 | +		ChatClientPromptContentObservationHandler customChatClientPromptContentObservationHandler() {  | 
 | 148 | +			return new ChatClientPromptContentObservationHandler();  | 
 | 149 | +		}  | 
 | 150 | + | 
 | 151 | +	}  | 
 | 152 | + | 
 | 153 | +	@Configuration(proxyBeanMethods = false)  | 
 | 154 | +	static class CustomTracingAwareLoggingObservationHandlerConfiguration {  | 
 | 155 | + | 
 | 156 | +		static TracingAwareLoggingObservationHandler<ChatClientObservationContext> handlerInstance = new TracingAwareLoggingObservationHandler<>(  | 
 | 157 | +				new ChatClientPromptContentObservationHandler(), null);  | 
 | 158 | + | 
 | 159 | +		@Bean  | 
 | 160 | +		TracingAwareLoggingObservationHandler<ChatClientObservationContext> chatClientPromptContentObservationHandler() {  | 
 | 161 | +			return handlerInstance;  | 
 | 162 | +		}  | 
 | 163 | + | 
51 | 164 | 	}  | 
52 | 165 | 
 
  | 
53 | 166 | }  | 
0 commit comments