Scale the number of test accounts used for SQL tests to one per shard. (#2468)

* Scale the number of test accounts used for SQL tests to one per shard.

* Set PLAYWRIGHT_SHARD_INDEX environment variable in CI workflow.

* Add log statement for the shared index and selected account.

* Remove console log.

* Fix order of accounts so that shard index maps to same account index.

* Try to fix the SQL account scope in ci.yml

* Get tokens for all accounts and use the shard index to pick which one.

* Set tokens without loop.

* Handcode the token use in tests.

* Fix database creation.

* Add debug for rbac token issues.

* Common function for retrieving NoSQL token.

* Disable eslint rule for noconsole temporarily.

* Move getNoSqlRbacToken to separate file.

* Fix ref to new function.

* mock Resource Graph API — fires on auto-subscription selection to populate account dropdown

* Code tidy-up.

* Fix build errors.

---------

Co-authored-by: Bikram Choudhury <bchoudhury@microsoft.com>

test/NoSqlTestSetup.ts

@@ -0,0 +1,71 @@
+export function getNoSqlRbacToken(): string | undefined {
+  let nosqlRbacToken: string | undefined;
+  const shardIndex = process.env.PLAYWRIGHT_SHARD_INDEX ?? "";
+  switch (parseInt(shardIndex)) {
+    case 1:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_1_TOKEN;
+      break;
+    case 2:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_2_TOKEN;
+      break;
+    case 3:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_3_TOKEN;
+      break;
+    case 4:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_4_TOKEN;
+      break;
+    case 5:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_5_TOKEN;
+      break;
+    case 6:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_6_TOKEN;
+      break;
+    case 7:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_7_TOKEN;
+      break;
+    case 8:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_8_TOKEN;
+      break;
+    case 9:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_9_TOKEN;
+      break;
+    case 10:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_10_TOKEN;
+      break;
+    case 11:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_11_TOKEN;
+      break;
+    case 12:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_12_TOKEN;
+      break;
+    case 13:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_13_TOKEN;
+      break;
+    case 14:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_14_TOKEN;
+      break;
+    case 15:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_15_TOKEN;
+      break;
+    case 16:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_16_TOKEN;
+      break;
+    case 17:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_17_TOKEN;
+      break;
+    case 18:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_18_TOKEN;
+      break;
+    case 19:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_19_TOKEN;
+      break;
+    case 20:
+      nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_20_TOKEN;
+      break;
+  }
+
+  if (!nosqlRbacToken) {
+    console.warn(`No NoSQL RBAC token found for shard index ${shardIndex}`);
+  }
+  return nosqlRbacToken;
+}

test/fx.ts

@@ -1,6 +1,7 @@
 import { DefaultAzureCredential } from "@azure/identity";
 import { Frame, Locator, Page, expect } from "@playwright/test";
 import crypto from "crypto";
+import { getNoSqlRbacToken } from "./NoSqlTestSetup";
 import { TestContainerContext } from "./testData";
 
 const RETRY_COUNT = 3;
@@ -43,17 +44,35 @@ export enum TestAccount {
   SQLContainerCopyOnly = "SQLContainerCopyOnly",
 }
 
-export const defaultAccounts: Record<TestAccount, string> = {
-  [TestAccount.Tables]: "github-e2etests-tables",
-  [TestAccount.Cassandra]: "github-e2etests-cassandra",
-  [TestAccount.Gremlin]: "github-e2etests-gremlin",
-  [TestAccount.Mongo]: "github-e2etests-mongo",
-  [TestAccount.MongoReadonly]: "github-e2etests-mongo-readonly",
-  [TestAccount.Mongo32]: "github-e2etests-mongo32",
-  [TestAccount.SQL]: "github-e2etests-sql",
-  [TestAccount.SQLReadOnly]: "github-e2etests-sql-readonly",
-  [TestAccount.SQLContainerCopyOnly]: "github-e2etests-sql-containercopyonly",
-};
+export function getDefaultAccountName(accountType: TestAccount): string {
+  switch (accountType) {
+    case TestAccount.Tables:
+      return "github-e2etests-tables";
+    case TestAccount.Cassandra:
+      return "github-e2etests-cassandra";
+    case TestAccount.Gremlin:
+      return "github-e2etests-gremlin";
+    case TestAccount.Mongo:
+      return "github-e2etests-mongo";
+    case TestAccount.MongoReadonly:
+      return "github-e2etests-mongo-readonly";
+    case TestAccount.Mongo32:
+      return "github-e2etests-mongo32";
+    case TestAccount.SQLReadOnly:
+      return "github-e2etests-sql-readonly";
+    case TestAccount.SQLContainerCopyOnly:
+      return "github-e2etests-sql-containercopyonly";
+    case TestAccount.SQL: {
+      const shardIndex = process.env.PLAYWRIGHT_SHARD_INDEX ?? "";
+      if (!shardIndex) {
+        throw new Error("PLAYWRIGHT_SHARD_INDEX is not set");
+      }
+      return "github-e2etests-sql-" + shardIndex;
+    }
+    default:
+      throw new Error(`No default account name defined for account type ${accountType}`);
+  }
+}
 
 export const resourceGroupName = process.env.DE_TEST_RESOURCE_GROUP ?? "de-e2e-tests";
 export const subscriptionId = process.env.DE_TEST_SUBSCRIPTION_ID ?? "69e02f2d-f059-4409-9eac-97e8a276ae2c";
@@ -77,7 +96,7 @@ export function getAccountName(accountType: TestAccount) {
   return (
     process.env[`DE_TEST_ACCOUNT_NAME_${accountType.toLocaleUpperCase()}`] ??
     tryGetStandardName(accountType) ??
-    defaultAccounts[accountType]
+    getDefaultAccountName(accountType)
   );
 }
 
@@ -102,7 +121,11 @@ export async function getTestExplorerUrl(accountType: TestAccount, options?: Tes
   // For now, since we don't test copilot, we can disable the copilot APIs by setting the feature flag to false.
   params.set("feature.enableCopilot", "false");
 
-  const nosqlRbacToken = process.env.NOSQL_TESTACCOUNT_TOKEN;
+  const nosqlRbacToken = getNoSqlRbacToken();
+  if (!nosqlRbacToken) {
+    throw new Error("No NOSQL RBAC token found.");
+  }
+
   const nosqlReadOnlyRbacToken = process.env.NOSQL_READONLY_TESTACCOUNT_TOKEN;
   const nosqlContainerCopyRbacToken = process.env.NOSQL_CONTAINERCOPY_TESTACCOUNT_TOKEN;
   const tableRbacToken = process.env.TABLE_TESTACCOUNT_TOKEN;

test/sql/containercopy/permissionsScreen.spec.ts

@@ -35,6 +35,36 @@ test.describe("Container Copy - Permission Screen Verification", () => {
     await expect(wrapper.getByTestId("CommandBar/Button:Refresh")).toBeVisible();
     await expect(wrapper.getByTestId("CommandBar/Button:Feedback")).toBeVisible();
 
+    // Mock Resource Graph API — fires on auto-subscription selection to populate account dropdown
+    await page.route(
+      "https://management.azure.com/providers/Microsoft.ResourceGraph/resources?api-version=2021-03-01",
+      async (route) => {
+        const request = route.request();
+        if (
+          request.method() === "POST" &&
+          (request.postDataJSON()?.query as string) ===
+            "resources | where type =~ 'microsoft.documentdb/databaseaccounts'"
+        ) {
+          const response = await route.fetch();
+          const responseData = await response.json();
+          if (responseData.data && Array.isArray(responseData.data)) {
+            responseData.data = responseData.data.map((d: any) => {
+              d.properties.backupPolicy.type = "Periodic";
+              return d;
+            });
+          }
+          await route.fulfill({
+            status: 200,
+            contentType: "application/json",
+            body: JSON.stringify(responseData),
+          });
+        } else {
+          await route.continue();
+        }
+      },
+      { times: 2 },
+    );
+
     // Open the Create Copy Job panel
     await createCopyJobButton.click();
     panel = frame.getByTestId("Panel:Create copy job");

test/sql/resourceToken.spec.ts

@@ -11,9 +11,10 @@ import {
   resourceGroupName,
   subscriptionId,
 } from "../fx";
+import { getNoSqlRbacToken } from "../NoSqlTestSetup";
 
 test("SQL account using Resource token", async ({ page }) => {
-  const nosqlAccountRbacToken = process.env.NOSQL_TESTACCOUNT_TOKEN || "";
+  const nosqlAccountRbacToken = getNoSqlRbacToken() ?? "";
   test.skip(nosqlAccountRbacToken.length > 0, "Resource tokens not supported when using data plane RBAC.");
 
   const credentials = getAzureCLICredentials();

test/testData.ts

@@ -18,6 +18,7 @@ import {
   subscriptionId,
   TestAccount,
 } from "./fx";
+import { getNoSqlRbacToken } from "./NoSqlTestSetup";
 
 // In Node.js >= 19, globalThis.crypto is already available as a read-only getter.
 // Only assign the polyfill for older versions.
@@ -134,7 +135,7 @@ async function createCosmosClientForSQLAccount(
 
   const rbacToken =
     accountType === TestAccount.SQL
-      ? process.env.NOSQL_TESTACCOUNT_TOKEN
+      ? getNoSqlRbacToken()
       : accountType === TestAccount.SQLContainerCopyOnly
       ? process.env.NOSQL_CONTAINERCOPY_TESTACCOUNT_TOKEN
       : "";

test/testExplorer/TestExplorer.ts

@@ -3,6 +3,7 @@ import "../../less/hostedexplorer.less";
 import { DataExplorerInputsFrame } from "../../src/Contracts/ViewModels";
 import { updateUserContext } from "../../src/UserContext";
 import { get, listKeys } from "../../src/Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { getNoSqlRbacToken } from "../NoSqlTestSetup";
 
 const urlSearchParams = new URLSearchParams(window.location.search);
 const resourceGroup = urlSearchParams.get("resourceGroup") || process.env.RESOURCE_GROUP || "";
@@ -15,8 +16,9 @@ const enablecontainercopy = urlSearchParams.get("enablecontainercopy");
 
 const nosqlRbacToken =
   urlSearchParams.get("nosqlRbacToken") ||
-  (enablecontainercopy ? process.env.NOSQL_CONTAINERCOPY_TESTACCOUNT_TOKEN : process.env.NOSQL_TESTACCOUNT_TOKEN) ||
+  (enablecontainercopy ? process.env.NOSQL_CONTAINERCOPY_TESTACCOUNT_TOKEN : getNoSqlRbacToken()) ||
   "";
+
 const nosqlReadOnlyRbacToken =
   urlSearchParams.get("nosqlReadOnlyRbacToken") || process.env.NOSQL_READONLY_TESTACCOUNT_TOKEN || "";
 const tableRbacToken = urlSearchParams.get("tableRbacToken") || process.env.TABLE_TESTACCOUNT_TOKEN || "";
@@ -70,6 +72,8 @@ const initTestExplorer = async (): Promise<void> => {
     updateUserContext({
       dataPlaneRbacEnabled: true,
     });
+  } else {
+    console.error(`No RBAC token found for test account type ${testAccountType}`);
   }
 
   const keys = await listKeys(subscriptionId, resourceGroup, accountName);