From 0d32d6ac639a7986f00d6a10ca6dd173c1765570 Mon Sep 17 00:00:00 2001
From: zhiyuanliang <zhiyuanliang@microsoft.com>
Date: Thu, 28 Aug 2025 23:26:52 +0800
Subject: [PATCH] update load balancing test

---
 test/loadBalance.test.ts | 15 ++++++++++-----
 test/utils/testHelper.ts |  9 ++++-----
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/test/loadBalance.test.ts b/test/loadBalance.test.ts
index b0b3c054..725bf525 100644
--- a/test/loadBalance.test.ts
+++ b/test/loadBalance.test.ts
@@ -6,10 +6,15 @@ import * as chaiAsPromised from "chai-as-promised";
 chai.use(chaiAsPromised);
 const expect = chai.expect;
 import { load } from "./exportedApi.js";
-import { MAX_TIME_OUT, restoreMocks, createMockedConnectionString, sleepInMs, createMockedEndpoint, mockConfigurationManagerGetClients, mockAppConfigurationClientLoadBalanceMode } from "./utils/testHelper.js";
+import { MAX_TIME_OUT, restoreMocks, createMockedConnectionString, createMockedKeyValue, sleepInMs, createMockedEndpoint, mockConfigurationManagerGetClients, mockAppConfigurationClientLoadBalanceMode } from "./utils/testHelper.js";
 import { AppConfigurationClient } from "@azure/app-configuration";
 import { ConfigurationClientWrapper } from "../src/configurationClientWrapper.js";
 
+const mockedKVs = [
+    { value: "red", key: "app.settings.fontColor" },
+    { value: "40", key: "app.settings.fontSize" },
+    { value: "30", key: "app.settings.fontSize", label: "prod" }
+].map(createMockedKeyValue);
 const fakeEndpoint_1 = createMockedEndpoint("fake_1");
 const fakeEndpoint_2 = createMockedEndpoint("fake_2");
 const fakeClientWrapper_1 = new ConfigurationClientWrapper(fakeEndpoint_1, new AppConfigurationClient(createMockedConnectionString(fakeEndpoint_1)));
@@ -29,8 +34,8 @@ describe("load balance", function () {
 
     it("should load balance the request when loadBalancingEnabled", async () => {
         mockConfigurationManagerGetClients([fakeClientWrapper_1, fakeClientWrapper_2], false);
-        mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_1, clientRequestCounter_1);
-        mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_2, clientRequestCounter_2);
+        mockAppConfigurationClientLoadBalanceMode([mockedKVs], fakeClientWrapper_1, clientRequestCounter_1);
+        mockAppConfigurationClientLoadBalanceMode([mockedKVs], fakeClientWrapper_2, clientRequestCounter_2);
 
         const connectionString = createMockedConnectionString();
         const settings = await load(connectionString, {
@@ -66,8 +71,8 @@ describe("load balance", function () {
         clientRequestCounter_1.count = 0;
         clientRequestCounter_2.count = 0;
         mockConfigurationManagerGetClients([fakeClientWrapper_1, fakeClientWrapper_2], false);
-        mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_1, clientRequestCounter_1);
-        mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_2, clientRequestCounter_2);
+        mockAppConfigurationClientLoadBalanceMode([mockedKVs], fakeClientWrapper_1, clientRequestCounter_1);
+        mockAppConfigurationClientLoadBalanceMode([mockedKVs], fakeClientWrapper_2, clientRequestCounter_2);
 
         const connectionString = createMockedConnectionString();
         // loadBalancingEnabled is default to false
diff --git a/test/utils/testHelper.ts b/test/utils/testHelper.ts
index 7e972ef9..9d9cc93e 100644
--- a/test/utils/testHelper.ts
+++ b/test/utils/testHelper.ts
@@ -121,12 +121,11 @@ function mockAppConfigurationClientListConfigurationSettings(pages: Configuratio
     });
 }
 
-function mockAppConfigurationClientLoadBalanceMode(clientWrapper: ConfigurationClientWrapper, countObject: { count: number }) {
-    const emptyPages: ConfigurationSetting[][] = [];
+function mockAppConfigurationClientLoadBalanceMode(pages: ConfigurationSetting[][], clientWrapper: ConfigurationClientWrapper, countObject: { count: number }) {
     sinon.stub(clientWrapper.client, "listConfigurationSettings").callsFake((listOptions) => {
         countObject.count += 1;
-        const kvs = _filterKVs(emptyPages.flat(), listOptions);
-        return getMockedIterator(emptyPages, kvs, listOptions);
+        const kvs = _filterKVs(pages.flat(), listOptions);
+        return getMockedIterator(pages, kvs, listOptions);
     });
 }
 
@@ -159,7 +158,7 @@ function mockConfigurationManagerGetClients(fakeClientWrappers: ConfigurationCli
     });
 }
 
-function mockAppConfigurationClientGetConfigurationSetting(kvList, customCallback?: (options) => any) {
+function mockAppConfigurationClientGetConfigurationSetting(kvList: any[], customCallback?: (options) => any) {
     sinon.stub(AppConfigurationClient.prototype, "getConfigurationSetting").callsFake((settingId, options) => {
         if (customCallback) {
             customCallback(options);