From 3600fa7cb97d62782fbfd7626370390ec14f5cca Mon Sep 17 00:00:00 2001
From: yutao <yutao.tao@bytedance.com>
Date: Fri, 27 Dec 2024 22:16:32 +0800
Subject: [PATCH] feat: add test case for azure

---
 connectivity-test/package.json               |  1 +
 connectivity-test/tests/connectivity.test.ts | 38 ++++++++++++++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/connectivity-test/package.json b/connectivity-test/package.json
index 91480df..6f8ccef 100644
--- a/connectivity-test/package.json
+++ b/connectivity-test/package.json
@@ -10,6 +10,7 @@
     "author": "",
     "license": "MIT",
     "devDependencies": {
+        "@azure/identity": "4.5.0",
         "@midscene/core": "latest",
         "@midscene/shared": "latest",
         "dotenv": "^16.4.5",
diff --git a/connectivity-test/tests/connectivity.test.ts b/connectivity-test/tests/connectivity.test.ts
index ae554f3..894e01d 100644
--- a/connectivity-test/tests/connectivity.test.ts
+++ b/connectivity-test/tests/connectivity.test.ts
@@ -1,9 +1,13 @@
 import { describe, it, expect, vi } from "vitest";
 import dotenv from "dotenv";
-import OpenAI from "openai";
+import OpenAI, { AzureOpenAI } from "openai";
 import { join } from "node:path";
 import { base64Encoded } from "@midscene/shared/img";
 import { callToGetJSONObject } from "@midscene/core/ai-model";
+import {
+  DefaultAzureCredential,
+  getBearerTokenProvider,
+} from "@azure/identity";
 
 // read and parse .env file
 const result = dotenv.config({
@@ -24,8 +28,8 @@ vi.setConfig({
 const imagePath = join(__dirname, "some_logo.png");
 const imageBase64 = base64Encoded(imagePath);
 
+const model = process.env.MIDSCENE_MODEL_NAME || "gpt-4o";
 describe("Use OpenAI SDK directly", () => {
-  const model = process.env.MIDSCENE_MODEL_NAME || "gpt-4o";
   it(`basic call with ${model}`, async () => {
     const openai = new OpenAI({
       apiKey: process.env.OPENAI_API_KEY,
@@ -94,3 +98,33 @@ describe("Use Midscene wrapped OpenAI SDK", () => {
     expect(result.content.content.length).toBeGreaterThan(5);
   });
 });
+
+// remove the ".skip" if you want to test Azure OpenAI Service
+describe.skip("Azure OpenAI Service", () => {
+  it("basic call", async () => {
+    // sample code: https://github.com/Azure/azure-sdk-for-js/blob/main/sdk/openai/openai/samples/cookbook/simpleCompletionsPage/app.js
+    const scope = process.env.MIDSCENE_AZURE_OPENAI_SCOPE;
+    if (typeof scope !== "string") {
+      throw new Error("MIDSCENE_AZURE_OPENAI_SCOPE is required");
+    }
+
+    const credential = new DefaultAzureCredential();
+    const tokenProvider = getBearerTokenProvider(credential, scope);
+
+    const extraAzureConfig = JSON.parse(
+      process.env.MIDSCENE_AZURE_OPENAI_INIT_CONFIG_JSON || "{}"
+    );
+    // console.log(extraAzureConfig);
+    const openai = new AzureOpenAI({
+      azureADTokenProvider: tokenProvider,
+      ...extraAzureConfig,
+    });
+
+    const response = await openai.chat.completions.create({
+      model: model,
+      messages: [{ role: "user", content: "Hello, how are you?" }],
+    });
+
+    expect(response.choices[0].message.content).toBeTruthy();
+  });
+});