Skip to content

Commit 35e6113

Browse files
sobychackomarkpollack
authored andcommitted
Adding integration test for Azure custom headers
1 parent b468354 commit 35e6113

File tree

1 file changed

+24
-16
lines changed

1 file changed

+24
-16
lines changed

spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/azure/AzureOpenAiAutoConfigurationIT.java

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
* See the License for the specific language governing permissions and
1414
* limitations under the License.
1515
*/
16+
1617
package org.springframework.ai.autoconfigure.azure;
1718

1819
import com.azure.ai.openai.OpenAIClient;
@@ -95,22 +96,29 @@ void chatCompletion() {
9596
}
9697

9798
@Test
98-
void httpRequestContainsUserAgentHeader() {
99-
contextRunner.run(context -> {
100-
OpenAIClient openAIClient = context.getBean(OpenAIClient.class);
101-
Field serviceClientField = ReflectionUtils.findField(OpenAIClient.class, "serviceClient");
102-
assertThat(serviceClientField).isNotNull();
103-
ReflectionUtils.makeAccessible(serviceClientField);
104-
OpenAIClientImpl oaci = (OpenAIClientImpl) ReflectionUtils.getField(serviceClientField, openAIClient);
105-
assertThat(oaci).isNotNull();
106-
HttpPipeline httpPipeline = oaci.getHttpPipeline();
107-
HttpResponse httpResponse = httpPipeline
108-
.send(new HttpRequest(HttpMethod.POST, new URI(System.getenv("AZURE_OPENAI_ENDPOINT")).toURL()))
109-
.block();
110-
assertThat(httpResponse).isNotNull();
111-
HttpHeader httpHeader = httpResponse.getRequest().getHeaders().get(HttpHeaderName.USER_AGENT);
112-
assertThat(httpHeader.getValue().startsWith("spring-ai azsdk-java-azure-ai-openai/")).isTrue();
113-
});
99+
void httpRequestContainsUserAgentAndCustomHeaders() {
100+
contextRunner
101+
.withPropertyValues("spring.ai.azure.openai.custom-headers.foo=bar",
102+
"spring.ai.azure.openai.custom-headers.fizz=buzz")
103+
.run(context -> {
104+
OpenAIClient openAIClient = context.getBean(OpenAIClient.class);
105+
Field serviceClientField = ReflectionUtils.findField(OpenAIClient.class, "serviceClient");
106+
assertThat(serviceClientField).isNotNull();
107+
ReflectionUtils.makeAccessible(serviceClientField);
108+
OpenAIClientImpl oaci = (OpenAIClientImpl) ReflectionUtils.getField(serviceClientField, openAIClient);
109+
assertThat(oaci).isNotNull();
110+
HttpPipeline httpPipeline = oaci.getHttpPipeline();
111+
HttpResponse httpResponse = httpPipeline
112+
.send(new HttpRequest(HttpMethod.POST, new URI(System.getenv("AZURE_OPENAI_ENDPOINT")).toURL()))
113+
.block();
114+
assertThat(httpResponse).isNotNull();
115+
HttpHeader httpHeader = httpResponse.getRequest().getHeaders().get(HttpHeaderName.USER_AGENT);
116+
assertThat(httpHeader.getValue().startsWith("spring-ai azsdk-java-azure-ai-openai/")).isTrue();
117+
HttpHeader customHeader1 = httpResponse.getRequest().getHeaders().get("foo");
118+
assertThat(customHeader1.getValue()).isEqualTo("bar");
119+
HttpHeader customHeader2 = httpResponse.getRequest().getHeaders().get("fizz");
120+
assertThat(customHeader2.getValue()).isEqualTo("buzz");
121+
});
114122
}
115123

116124
@Test

0 commit comments

Comments
 (0)