Initial plan

src/Contracts/DataModels.ts

@@ -275,7 +275,8 @@ export interface DataMaskingPolicy {
     startPosition: number;
     length: number;
   }>;
-  excludedPaths?: string[];
+  excludedPaths: string[];
+  isPolicyEnabled: boolean;
 }
 
 export interface MaterializedView {

src/Explorer/Controls/Settings/SettingsComponent.test.tsx

@@ -30,6 +30,7 @@ jest.mock("../../../Common/dataAccess/updateCollection", () => ({
     dataMaskingPolicy: {
       includedPaths: [],
       excludedPaths: ["/excludedPath"],
+      isPolicyEnabled: true,
     },
     indexes: [],
   }),
@@ -306,10 +307,12 @@ describe("SettingsComponent", () => {
       dataMaskingContent: {
         includedPaths: [],
         excludedPaths: ["/excludedPath"],
+        isPolicyEnabled: true,
       },
       dataMaskingContentBaseline: {
         includedPaths: [],
         excludedPaths: [],
+        isPolicyEnabled: false,
       },
       isDataMaskingDirty: true,
     });
@@ -323,6 +326,7 @@ describe("SettingsComponent", () => {
     expect(wrapper.state("dataMaskingContentBaseline")).toEqual({
       includedPaths: [],
       excludedPaths: ["/excludedPath"],
+      isPolicyEnabled: true,
     });
   });
 
@@ -336,6 +340,7 @@ describe("SettingsComponent", () => {
     const invalidPolicy: InvalidPolicy = {
       includedPaths: "invalid",
       excludedPaths: [],
+      isPolicyEnabled: false,
     };
     // Use type assertion since we're deliberately testing with invalid data
     settingsComponentInstance["onDataMaskingContentChange"](invalidPolicy as unknown as DataModels.DataMaskingPolicy);
@@ -344,6 +349,7 @@ describe("SettingsComponent", () => {
     expect(wrapper.state("dataMaskingContent")).toEqual({
       includedPaths: "invalid",
       excludedPaths: [],
+      isPolicyEnabled: false,
     });
     expect(wrapper.state("dataMaskingValidationErrors")).toEqual(["includedPaths must be an array"]);
 
@@ -358,6 +364,7 @@ describe("SettingsComponent", () => {
         },
       ],
       excludedPaths: ["/excludedPath"],
+      isPolicyEnabled: true,
     };
 
     settingsComponentInstance["onDataMaskingContentChange"](validPolicy);
@@ -381,6 +388,7 @@ describe("SettingsComponent", () => {
         },
       ],
       excludedPaths: ["/excludedPath1"],
+      isPolicyEnabled: false,
     };
 
     const modifiedPolicy = {
@@ -393,6 +401,7 @@ describe("SettingsComponent", () => {
         },
       ],
       excludedPaths: ["/excludedPath2"],
+      isPolicyEnabled: true,
     };
 
     // Set initial state

src/Explorer/Controls/Settings/SettingsComponent.tsx

@@ -687,14 +687,22 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
     this.setState({ isComputedPropertiesDirty: isComputedPropertiesDirty });
 
   private onDataMaskingContentChange = (newDataMasking: DataModels.DataMaskingPolicy): void => {
+    if (!newDataMasking.excludedPaths) {
+      newDataMasking.excludedPaths = [];
+    }
+    if (!newDataMasking.includedPaths) {
+      newDataMasking.includedPaths = [];
+    }
+
     const validationErrors = [];
-    if (newDataMasking.includedPaths === undefined || newDataMasking.includedPaths === null) {
-      validationErrors.push("includedPaths is required");
-    } else if (!Array.isArray(newDataMasking.includedPaths)) {
+    if (!Array.isArray(newDataMasking.includedPaths)) {
       validationErrors.push("includedPaths must be an array");
     }
-    if (newDataMasking.excludedPaths !== undefined && !Array.isArray(newDataMasking.excludedPaths)) {
-      validationErrors.push("excludedPaths must be an array if provided");
+    if (!Array.isArray(newDataMasking.excludedPaths)) {
+      validationErrors.push("excludedPaths must be an array");
+    }
+    if (typeof newDataMasking.isPolicyEnabled !== "boolean") {
+      validationErrors.push("isPolicyEnabled must be a boolean");
     }
 
     this.setState({
@@ -835,6 +843,7 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
     const dataMaskingContent: DataModels.DataMaskingPolicy = {
       includedPaths: this.collection.dataMaskingPolicy?.()?.includedPaths || [],
       excludedPaths: this.collection.dataMaskingPolicy?.()?.excludedPaths || [],
+      isPolicyEnabled: this.collection.dataMaskingPolicy?.()?.isPolicyEnabled ?? true,
     };
     const conflictResolutionPolicy: DataModels.ConflictResolutionPolicy =
       this.collection.conflictResolutionPolicy && this.collection.conflictResolutionPolicy();

src/Explorer/Controls/Settings/SettingsSubComponents/DataMaskingComponent.test.tsx

@@ -53,6 +53,7 @@ describe("DataMaskingComponent", () => {
       },
     ],
     excludedPaths: [],
+    isPolicyEnabled: false,
   };
 
   let changeContentCallback: () => void;
@@ -77,7 +78,7 @@ describe("DataMaskingComponent", () => {
       <DataMaskingComponent
         {...mockProps}
         dataMaskingContent={samplePolicy}
-        dataMaskingContentBaseline={{ ...samplePolicy, excludedPaths: ["/excluded"] }}
+        dataMaskingContentBaseline={{ ...samplePolicy, isPolicyEnabled: true }}
       />,
     );
 
@@ -122,7 +123,7 @@ describe("DataMaskingComponent", () => {
   });
 
   it("resets content when shouldDiscardDataMasking is true", async () => {
-    const baselinePolicy = { ...samplePolicy, excludedPaths: ["/excluded"] };
+    const baselinePolicy = { ...samplePolicy, isPolicyEnabled: true };
 
     const wrapper = mount(
       <DataMaskingComponent
@@ -158,7 +159,7 @@ describe("DataMaskingComponent", () => {
     wrapper.update();
 
     // Update baseline to trigger componentDidUpdate
-    const newBaseline = { ...samplePolicy, excludedPaths: ["/excluded"] };
+    const newBaseline = { ...samplePolicy, isPolicyEnabled: true };
     wrapper.setProps({ dataMaskingContentBaseline: newBaseline });
 
     expect(mockProps.onDataMaskingDirtyChange).toHaveBeenCalledWith(true);
@@ -173,6 +174,7 @@ describe("DataMaskingComponent", () => {
     const invalidPolicy: Record<string, unknown> = {
       includedPaths: "not an array",
       excludedPaths: [] as string[],
+      isPolicyEnabled: "not a boolean",
     };
 
     mockGetValue.mockReturnValue(JSON.stringify(invalidPolicy));
@@ -195,7 +197,7 @@ describe("DataMaskingComponent", () => {
     wrapper.update();
 
     // First change
-    const modifiedPolicy1 = { ...samplePolicy, excludedPaths: ["/path1"] };
+    const modifiedPolicy1 = { ...samplePolicy, isPolicyEnabled: true };
     mockGetValue.mockReturnValue(JSON.stringify(modifiedPolicy1));
     changeContentCallback();
     expect(mockProps.onDataMaskingDirtyChange).toHaveBeenCalledWith(true);

src/Explorer/Controls/Settings/SettingsSubComponents/DataMaskingComponent.tsx

@@ -4,7 +4,7 @@ import * as React from "react";
 import * as DataModels from "../../../../Contracts/DataModels";
 import { loadMonaco } from "../../../LazyMonaco";
 import { titleAndInputStackProps, unsavedEditorWarningMessage } from "../SettingsRenderUtils";
-import { isDirty as isContentDirty, isDataMaskingEnabled } from "../SettingsUtils";
+import { isDataMaskingEnabled, isDirty as isContentDirty } from "../SettingsUtils";
 
 export interface DataMaskingComponentProps {
   shouldDiscardDataMasking: boolean;
@@ -22,8 +22,16 @@ interface DataMaskingComponentState {
 }
 
 const emptyDataMaskingPolicy: DataModels.DataMaskingPolicy = {
-  includedPaths: [],
+  includedPaths: [
+    {
+      path: "/",
+      strategy: "Default",
+      startPosition: 0,
+      length: -1,
+    },
+  ],
   excludedPaths: [],
+  isPolicyEnabled: true,
 };
 
 export class DataMaskingComponent extends React.Component<DataMaskingComponentProps, DataMaskingComponentState> {

src/Explorer/Controls/Settings/TestUtils.tsx

@@ -68,6 +68,7 @@ export const collection = {
   dataMaskingPolicy: ko.observable<DataModels.DataMaskingPolicy>({
     includedPaths: [],
     excludedPaths: ["/excludedPath"],
+    isPolicyEnabled: true,
   }),
   readSettings: () => {
     return;

src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap

@@ -638,6 +638,7 @@ exports[`SettingsComponent renders 1`] = `
                   "/excludedPath",
                 ],
                 "includedPaths": [],
+                "isPolicyEnabled": true,
               }
             }
             dataMaskingContentBaseline={
@@ -646,6 +647,7 @@ exports[`SettingsComponent renders 1`] = `
                   "/excludedPath",
                 ],
                 "includedPaths": [],
+                "isPolicyEnabled": true,
               }
             }
             onDataMaskingContentChange={[Function]}

src/Explorer/Tree/Collection.ts

@@ -141,6 +141,7 @@ export default class Collection implements ViewModels.Collection {
     const defaultDataMaskingPolicy: DataModels.DataMaskingPolicy = {
       includedPaths: Array<{ path: string; strategy: string; startPosition: number; length: number }>(),
       excludedPaths: Array<string>(),
+      isPolicyEnabled: true,
     };
     const observablePolicy = ko.observable(data.dataMaskingPolicy || defaultDataMaskingPolicy);
     observablePolicy.subscribe(() => {});

test/fx.ts

@@ -58,9 +58,7 @@ export const defaultAccounts: Record<TestAccount, string> = {
 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";
 export const TEST_AUTOSCALE_THROUGHPUT_RU = 1000;
-export const TEST_MANUAL_THROUGHPUT_RU = 800;
 export const TEST_AUTOSCALE_MAX_THROUGHPUT_RU_2K = 2000;
-export const TEST_AUTOSCALE_MAX_THROUGHPUT_RU_4K = 4000;
 export const TEST_MANUAL_THROUGHPUT_RU_2K = 2000;
 export const ONE_MINUTE_MS: number = 60 * 1000;
 

test/sql/scaleAndSettings/sharedThroughput.spec.ts

@@ -1,229 +0,0 @@
-import { Locator, expect, test } from "@playwright/test";
-import {
-  CommandBarButton,
-  DataExplorer,
-  ONE_MINUTE_MS,
-  TEST_AUTOSCALE_MAX_THROUGHPUT_RU_4K,
-  TEST_MANUAL_THROUGHPUT_RU,
-  TestAccount,
-} from "../../fx";
-import { TestDatabaseContext, createTestDB } from "../../testData";
-
-test.describe("Database with Shared Throughput", () => {
-  let dbContext: TestDatabaseContext = null!;
-  let explorer: DataExplorer = null!;
-  const containerId = "sharedcontainer";
-
-  // Helper methods
-  const getThroughputInput = (type: "manual" | "autopilot"): Locator => {
-    return explorer.frame.getByTestId(`${type}-throughput-input`);
-  };
-
-  test.afterEach("Delete Test Database", async () => {
-    await dbContext?.dispose();
-  });
-
-  test.describe("Manual Throughput Tests", () => {
-    test.beforeEach(async ({ page }) => {
-      explorer = await DataExplorer.open(page, TestAccount.SQL);
-    });
-
-    test("Create database with shared manual throughput and verify Scale node in UI", async () => {
-      test.setTimeout(120000); // 2 minutes timeout
-      // Create database with shared manual throughput (400 RU/s)
-      dbContext = await createTestDB({ throughput: 400 });
-
-      // Verify database node appears in the tree
-      const databaseNode = await explorer.waitForNode(dbContext.database.id);
-      expect(databaseNode).toBeDefined();
-
-      // Expand the database node to see child nodes
-      await databaseNode.expand();
-
-      // Verify that "Scale" node appears under the database
-      const scaleNode = await explorer.waitForNode(`${dbContext.database.id}/Scale`);
-      expect(scaleNode).toBeDefined();
-      await expect(scaleNode.element).toBeVisible();
-    });
-
-    test("Add container to shared database without dedicated throughput", async () => {
-      // Create database with shared manual throughput
-      dbContext = await createTestDB({ throughput: 400 });
-
-      // Wait for the database to appear in the tree
-      await explorer.waitForNode(dbContext.database.id);
-
-      // Add a container to the shared database via UI
-      const newContainerButton = await explorer.globalCommandButton("New Container");
-      await newContainerButton.click();
-
-      await explorer.whilePanelOpen(
-        "New Container",
-        async (panel, okButton) => {
-          // Select "Use existing" database
-          const useExistingRadio = panel.getByRole("radio", { name: /Use existing/i });
-          await useExistingRadio.click();
-
-          // Select the database from dropdown using the new data-testid
-          const databaseDropdown = panel.getByRole("combobox", { name: "Choose an existing database" });
-          await databaseDropdown.click();
-
-          await explorer.frame.getByRole("option", { name: dbContext.database.id }).click();
-          // Now you can target the specific database option by its data-testid
-          //await panel.getByTestId(`database-option-${dbContext.database.id}`).click();
-          // Fill container id
-          await panel.getByRole("textbox", { name: "Container id, Example Container1" }).fill(containerId);
-
-          // Fill partition key
-          await panel.getByRole("textbox", { name: "Partition key" }).fill("/pk");
-
-          // Ensure "Provision dedicated throughput" is NOT checked
-          const dedicatedThroughputCheckbox = panel.getByRole("checkbox", {
-            name: /Provision dedicated throughput for this container/i,
-          });
-
-          if (await dedicatedThroughputCheckbox.isVisible()) {
-            const isChecked = await dedicatedThroughputCheckbox.isChecked();
-            if (isChecked) {
-              await dedicatedThroughputCheckbox.uncheck();
-            }
-          }
-
-          await okButton.click();
-        },
-        { closeTimeout: 5 * ONE_MINUTE_MS },
-      );
-
-      // Verify container was created under the database
-      const containerNode = await explorer.waitForContainerNode(dbContext.database.id, containerId);
-      expect(containerNode).toBeDefined();
-    });
-
-    test("Scale shared database manual throughput", async () => {
-      // Create database with shared manual throughput (400 RU/s)
-      dbContext = await createTestDB({ throughput: 400 });
-
-      // Navigate to the scale settings by clicking the "Scale" node in the tree
-      const databaseNode = await explorer.waitForNode(dbContext.database.id);
-      await databaseNode.expand();
-      const scaleNode = await explorer.waitForNode(`${dbContext.database.id}/Scale`);
-      await scaleNode.element.click();
-
-      // Update manual throughput from 400 to 800
-      await getThroughputInput("manual").fill(TEST_MANUAL_THROUGHPUT_RU.toString());
-
-      // Save changes
-      await explorer.commandBarButton(CommandBarButton.Save).click();
-
-      // Verify success message
-      await expect(explorer.getConsoleHeaderStatus()).toContainText(
-        `Successfully updated offer for database ${dbContext.database.id}`,
-        {
-          timeout: 2 * ONE_MINUTE_MS,
-        },
-      );
-    });
-
-    test("Scale shared database from manual to autoscale", async () => {
-      // Create database with shared manual throughput (400 RU/s)
-      dbContext = await createTestDB({ throughput: 400 });
-
-      // Open database settings by clicking the "Scale" node
-      const databaseNode = await explorer.waitForNode(dbContext.database.id);
-      await databaseNode.expand();
-      const scaleNode = await explorer.waitForNode(`${dbContext.database.id}/Scale`);
-      await scaleNode.element.click();
-
-      // Switch to Autoscale
-      const autoscaleRadio = explorer.frame.getByText("Autoscale", { exact: true });
-      await autoscaleRadio.click();
-
-      // Set autoscale max throughput to 1000
-      //await getThroughputInput("autopilot").fill(TEST_AUTOSCALE_THROUGHPUT_RU.toString());
-
-      // Save changes
-      await explorer.commandBarButton(CommandBarButton.Save).click();
-
-      await expect(explorer.getConsoleHeaderStatus()).toContainText(
-        `Successfully updated offer for database ${dbContext.database.id}`,
-        {
-          timeout: 2 * ONE_MINUTE_MS,
-        },
-      );
-    });
-  });
-
-  test.describe("Autoscale Throughput Tests", () => {
-    test.beforeEach(async ({ page }) => {
-      explorer = await DataExplorer.open(page, TestAccount.SQL);
-    });
-
-    test("Create database with shared autoscale throughput and verify Scale node in UI", async () => {
-      test.setTimeout(120000); // 2 minutes timeout
-
-      // Create database with shared autoscale throughput (max 1000 RU/s)
-      dbContext = await createTestDB({ maxThroughput: 1000 });
-
-      // Verify database node appears
-      const databaseNode = await explorer.waitForNode(dbContext.database.id);
-      expect(databaseNode).toBeDefined();
-
-      // Expand the database node to see child nodes
-      await databaseNode.expand();
-
-      // Verify that "Scale" node appears under the database
-      const scaleNode = await explorer.waitForNode(`${dbContext.database.id}/Scale`);
-      expect(scaleNode).toBeDefined();
-      await expect(scaleNode.element).toBeVisible();
-    });
-
-    test("Scale shared database autoscale throughput", async () => {
-      // Create database with shared autoscale throughput (max 1000 RU/s)
-      dbContext = await createTestDB({ maxThroughput: 1000 });
-
-      // Open database settings
-      const databaseNode = await explorer.waitForNode(dbContext.database.id);
-      await databaseNode.expand();
-      const scaleNode = await explorer.waitForNode(`${dbContext.database.id}/Scale`);
-      await scaleNode.element.click();
-
-      // Update autoscale max throughput from 1000 to 4000
-      await getThroughputInput("autopilot").fill(TEST_AUTOSCALE_MAX_THROUGHPUT_RU_4K.toString());
-
-      // Save changes
-      await explorer.commandBarButton(CommandBarButton.Save).click();
-
-      // Verify success message
-      await expect(explorer.getConsoleHeaderStatus()).toContainText(
-        `Successfully updated offer for database ${dbContext.database.id}`,
-        {
-          timeout: 2 * ONE_MINUTE_MS,
-        },
-      );
-    });
-
-    test("Scale shared database from autoscale to manual", async () => {
-      // Create database with shared autoscale throughput (max 1000 RU/s)
-      dbContext = await createTestDB({ maxThroughput: 1000 });
-
-      // Open database settings
-      const databaseNode = await explorer.waitForNode(dbContext.database.id);
-      await databaseNode.expand();
-      const scaleNode = await explorer.waitForNode(`${dbContext.database.id}/Scale`);
-      await scaleNode.element.click();
-
-      // Switch to Manual
-      const manualRadio = explorer.frame.getByText("Manual", { exact: true });
-      await manualRadio.click();
-
-      // Save changes
-      await explorer.commandBarButton(CommandBarButton.Save).click();
-
-      // Verify success message
-      await expect(explorer.getConsoleHeaderStatus()).toContainText(
-        `Successfully updated offer for database ${dbContext.database.id}`,
-        { timeout: 2 * ONE_MINUTE_MS },
-      );
-    });
-  });
-});

test/testData.ts

@@ -82,75 +82,6 @@ export class TestContainerContext {
   }
 }
 
-export class TestDatabaseContext {
-  constructor(
-    public armClient: CosmosDBManagementClient,
-    public client: CosmosClient,
-    public database: Database,
-  ) {}
-
-  async dispose() {
-    await this.database.delete();
-  }
-}
-
-export interface CreateTestDBOptions {
-  throughput?: number;
-  maxThroughput?: number; // For autoscale
-}
-
-// Helper function to create ARM client and Cosmos client for SQL account
-async function createCosmosClientForSQLAccount(
-  accountType: TestAccount.SQL | TestAccount.SQLContainerCopyOnly = TestAccount.SQL,
-): Promise<{ armClient: CosmosDBManagementClient; client: CosmosClient }> {
-  const credentials = getAzureCLICredentials();
-  const adaptedCredentials = new AzureIdentityCredentialAdapter(credentials);
-  const armClient = new CosmosDBManagementClient(adaptedCredentials, subscriptionId);
-  const accountName = getAccountName(accountType);
-  const account = await armClient.databaseAccounts.get(resourceGroupName, accountName);
-
-  const clientOptions: CosmosClientOptions = {
-    endpoint: account.documentEndpoint!,
-  };
-
-  const rbacToken =
-    accountType === TestAccount.SQL
-      ? process.env.NOSQL_TESTACCOUNT_TOKEN
-      : accountType === TestAccount.SQLContainerCopyOnly
-      ? process.env.NOSQL_CONTAINERCOPY_TESTACCOUNT_TOKEN
-      : "";
-
-  if (rbacToken) {
-    clientOptions.tokenProvider = async (): Promise<string> => {
-      const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
-      const authorizationToken = `${AUTH_PREFIX}${rbacToken}`;
-      return authorizationToken;
-    };
-  } else {
-    const keys = await armClient.databaseAccounts.listKeys(resourceGroupName, accountName);
-    clientOptions.key = keys.primaryMasterKey;
-  }
-
-  const client = new CosmosClient(clientOptions);
-
-  return { armClient, client };
-}
-
-export async function createTestDB(options?: CreateTestDBOptions): Promise<TestDatabaseContext> {
-  const databaseId = generateUniqueName("db");
-  const { armClient, client } = await createCosmosClientForSQLAccount();
-
-  // Create database with provisioned throughput (shared throughput)
-  // This checks the "Provision database throughput" option
-  const { database } = await client.databases.create({
-    id: databaseId,
-    throughput: options?.throughput, // Manual throughput (e.g., 400)
-    maxThroughput: options?.maxThroughput, // Autoscale max throughput (e.g., 1000)
-  });
-
-  return new TestDatabaseContext(armClient, client, database);
-}
-
 type createTestSqlContainerConfig = {
   includeTestData?: boolean;
   partitionKey?: string;
@@ -173,7 +104,34 @@ export async function createMultipleTestContainers({
   const creationPromises: Promise<TestContainerContext>[] = [];
 
   const databaseId = databaseName ? databaseName : generateUniqueName("db");
-  const { armClient, client } = await createCosmosClientForSQLAccount(accountType);
+  const credentials = getAzureCLICredentials();
+  const adaptedCredentials = new AzureIdentityCredentialAdapter(credentials);
+  const armClient = new CosmosDBManagementClient(adaptedCredentials, subscriptionId);
+  const accountName = getAccountName(accountType);
+  const account = await armClient.databaseAccounts.get(resourceGroupName, accountName);
+
+  const clientOptions: CosmosClientOptions = {
+    endpoint: account.documentEndpoint!,
+  };
+
+  const rbacToken =
+    accountType === TestAccount.SQL
+      ? process.env.NOSQL_TESTACCOUNT_TOKEN
+      : accountType === TestAccount.SQLContainerCopyOnly
+      ? process.env.NOSQL_CONTAINERCOPY_TESTACCOUNT_TOKEN
+      : "";
+  if (rbacToken) {
+    clientOptions.tokenProvider = async (): Promise<string> => {
+      const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
+      const authorizationToken = `${AUTH_PREFIX}${rbacToken}`;
+      return authorizationToken;
+    };
+  } else {
+    const keys = await armClient.databaseAccounts.listKeys(resourceGroupName, accountName);
+    clientOptions.key = keys.primaryMasterKey;
+  }
+
+  const client = new CosmosClient(clientOptions);
   const { database } = await client.databases.createIfNotExists({ id: databaseId });
 
   try {
@@ -200,8 +158,29 @@ export async function createTestSQLContainer({
 }: createTestSqlContainerConfig = {}) {
   const databaseId = databaseName ? databaseName : generateUniqueName("db");
   const containerId = "testcontainer"; // A unique container name isn't needed because the database is unique
-  const { armClient, client } = await createCosmosClientForSQLAccount();
+  const credentials = getAzureCLICredentials();
+  const adaptedCredentials = new AzureIdentityCredentialAdapter(credentials);
+  const armClient = new CosmosDBManagementClient(adaptedCredentials, subscriptionId);
+  const accountName = getAccountName(TestAccount.SQL);
+  const account = await armClient.databaseAccounts.get(resourceGroupName, accountName);
 
+  const clientOptions: CosmosClientOptions = {
+    endpoint: account.documentEndpoint!,
+  };
+
+  const nosqlAccountRbacToken = process.env.NOSQL_TESTACCOUNT_TOKEN;
+  if (nosqlAccountRbacToken) {
+    clientOptions.tokenProvider = async (): Promise<string> => {
+      const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
+      const authorizationToken = `${AUTH_PREFIX}${nosqlAccountRbacToken}`;
+      return authorizationToken;
+    };
+  } else {
+    const keys = await armClient.databaseAccounts.listKeys(resourceGroupName, accountName);
+    clientOptions.key = keys.primaryMasterKey;
+  }
+
+  const client = new CosmosClient(clientOptions);
   const { database } = await client.databases.createIfNotExists({ id: databaseId });
   try {
     const { container } = await database.containers.createIfNotExists({