Add sharedThroughput db tests (#2287)

* Add sharedThroughput db tests

* Run npm commands

* Run npm format

* Fix error with adding container to shared db test

* Add changes for Load more option to work

* Remove localhost

* Add Mongo Pagination tests

* Run npm format

* Revert "Add Mongo Pagination tests"

This reverts commit 50a244e6f9.

* Minor fixes

* Fix errors

* Minor fix

* Fix tests

* Address comments

* Address comments

.github/workflows/cleanup.yml

@@ -6,8 +6,8 @@ on:
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:
   schedule:
-    # Once every two hours
+    # Once every day at 7 AM PST
-    - cron: "0 */2 * * *"
+    - cron: "0 13 * * *"
 
 permissions:
   id-token: write

test/fx.ts

@@ -58,7 +58,9 @@ 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;
 
@@ -378,9 +380,11 @@ type PanelOpenOptions = {
 
 export enum CommandBarButton {
   Save = "Save",
+  Delete = "Delete",
   Execute = "Execute",
   ExecuteQuery = "Execute Query",
   UploadItem = "Upload Item",
+  NewDocument = "New Document",
 }
 
 /** Helper class that provides locator methods for DataExplorer components, on top of a Frame */
@@ -478,7 +482,7 @@ export class DataExplorer {
     return await this.waitForNode(`${databaseId}/${containerId}/Documents`);
   }
 
-  async waitForCommandBarButton(label: string, timeout?: number): Promise<Locator> {
+  async waitForCommandBarButton(label: CommandBarButton, timeout?: number): Promise<Locator> {
     const commandBar = this.commandBarButton(label);
     await commandBar.waitFor({ state: "visible", timeout });
     return commandBar;
@@ -515,15 +519,6 @@ export class DataExplorer {
     const containerNode = await this.waitForContainerNode(context.database.id, context.container.id);
     await containerNode.expand();
 
-    // refresh tree to remove deleted database
-    const consoleMessages = await this.getNotificationConsoleMessages();
-    const refreshButton = this.frame.getByTestId("Sidebar/RefreshButton");
-    await refreshButton.click();
-    await expect(consoleMessages).toContainText("Successfully refreshed databases", {
-      timeout: ONE_MINUTE_MS,
-    });
-    await this.collapseNotificationConsole();
-
     const scaleAndSettingsButton = this.frame.getByTestId(
       `TreeNode:${context.database.id}/${context.container.id}/Scale & Settings`,
     );

test/mongo/document.spec.ts

@@ -1,7 +1,7 @@
 import { expect, test } from "@playwright/test";
 
 import { setupCORSBypass } from "../CORSBypass";
-import { DataExplorer, DocumentsTab, TestAccount } from "../fx";
+import { CommandBarButton, DataExplorer, DocumentsTab, TestAccount } from "../fx";
 import { retry, serializeMongoToJson, setPartitionKeys } from "../testData";
 import { documentTestCases } from "./testCases";
 
@@ -48,19 +48,20 @@ for (const { name, databaseId, containerId, documents } of documentTestCases) {
           expect(resultData?._id).not.toBeNull();
           expect(resultData?._id).toEqual(docId);
         });
-        test(`should be able to create and delete new document from ${docId}`, async () => {
+        test(`should be able to create and delete new document from ${docId}`, async ({ page }) => {
           const span = documentsTab.documentsListPane.getByText(docId, { exact: true }).nth(0);
           await span.waitFor();
           await expect(span).toBeVisible();
 
           await span.click();
+          await page.waitForTimeout(5000); // wait for 5 seconds to ensure document is fully loaded. waitforTimeout is not recommended generally but here we are working around flakiness in the test env
+
           let newDocumentId;
           await retry(async () => {
-            const newDocumentButton = await explorer.waitForCommandBarButton("New Document", 5000);
+            const newDocumentButton = await explorer.waitForCommandBarButton(CommandBarButton.NewDocument, 5000);
             await expect(newDocumentButton).toBeVisible();
             await expect(newDocumentButton).toBeEnabled();
             await newDocumentButton.click();
-
             await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
 
             newDocumentId = `${Date.now().toString()}-delete`;
@@ -71,8 +72,9 @@ for (const { name, databaseId, containerId, documents } of documentTestCases) {
             };
 
             await documentsTab.resultsEditor.setText(JSON.stringify(newDocument));
-            const saveButton = await explorer.waitForCommandBarButton("Save", 5000);
+            const saveButton = await explorer.waitForCommandBarButton(CommandBarButton.Save, 5000);
             await saveButton.click({ timeout: 5000 });
+
             await expect(saveButton).toBeHidden({ timeout: 5000 });
           }, 3);
 
@@ -84,7 +86,7 @@ for (const { name, databaseId, containerId, documents } of documentTestCases) {
           await newSpan.click();
           await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
 
-          const deleteButton = await explorer.waitForCommandBarButton("Delete", 5000);
+          const deleteButton = await explorer.waitForCommandBarButton(CommandBarButton.Delete, 5000);
           await deleteButton.click();
 
           const deleteDialogButton = await explorer.waitForDialogButton("Delete", 5000);

test/sql/document.spec.ts

@@ -136,9 +136,7 @@ test.describe.serial("Upload Item", () => {
     if (existsSync(uploadDocumentDirPath)) {
       rmdirSync(uploadDocumentDirPath);
     }
-    if (!process.env.CI) {
     await context?.dispose();
-    }
   });
 
   test.afterEach("Close Upload Items panel if still open", async () => {

test/sql/query.spec.ts

@@ -30,12 +30,9 @@ test.beforeEach("Open new query tab", async ({ page }) => {
   await explorer.frame.getByTestId("NotificationConsole/Contents").waitFor();
 });
 
-// Delete database only if not running in CI
-if (!process.env.CI) {
 test.afterAll("Delete Test Database", async () => {
   await context?.dispose();
 });
-}
 
 test("Query results", async () => {
   // Run the query and verify the results

test/sql/scaleAndSettings/changePartitionKey.spec.ts

@@ -23,12 +23,9 @@ test.describe("Change Partition Key", () => {
     await PartitionKeyTab.click();
   });
 
-  // Delete database only if not running in CI
-  if (!process.env.CI) {
   test.afterEach("Delete Test Database", async () => {
     await context?.dispose();
   });
-  }
 
   test("Change partition key path", async ({ page }) => {
     await expect(explorer.frame.getByText("/partitionKey")).toBeVisible();

test/sql/scaleAndSettings/scale.spec.ts

@@ -118,7 +118,5 @@ async function openScaleTab(browser: Browser): Promise<SetupResult> {
 }
 
 async function cleanup({ context }: Partial<SetupResult>) {
-  if (!process.env.CI) {
   await context?.dispose();
 }
-}

test/sql/scaleAndSettings/settings.spec.ts

@@ -17,12 +17,9 @@ test.describe("Settings under Scale & Settings", () => {
     await settingsTab.click();
   });
 
-  // Delete database only if not running in CI
-  if (!process.env.CI) {
   test.afterAll("Delete Test Database", async () => {
     await context?.dispose();
   });
-  }
 
   test("Update TTL to On (no default)", async () => {
     const ttlOnNoDefaultRadioButton = explorer.frame.getByRole("radio", { name: "ttl-on-no-default-option" });

test/sql/scaleAndSettings/sharedThroughput.spec.ts

@@ -0,0 +1,229 @@
+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/sql/scripts/storedProcedure.spec.ts

@@ -43,7 +43,7 @@ test.describe("Stored Procedures", () => {
     );
 
     // Execute stored procedure
-    const executeButton = explorer.commandBarButton(CommandBarButton.Execute);
+    const executeButton = explorer.commandBarButton(CommandBarButton.Execute).first();
     await executeButton.click();
     const executeSidePanelButton = explorer.frame.getByTestId("Panel/OkButton");
     await executeSidePanelButton.click();

test/sql/scripts/trigger.spec.ts

@@ -26,11 +26,9 @@ test.describe("Triggers", () => {
     explorer = await DataExplorer.open(page, TestAccount.SQL);
   });
 
-  if (!process.env.CI) {
   test.afterAll("Delete Test Database", async () => {
     await context?.dispose();
   });
-  }
 
   test("Add and delete trigger", async ({ page }, testInfo) => {
     // Open container context menu and click New Trigger

test/sql/scripts/userDefinedFunction.spec.ts

@@ -19,11 +19,9 @@ test.describe("User Defined Functions", () => {
     explorer = await DataExplorer.open(page, TestAccount.SQL);
   });
 
-  if (!process.env.CI) {
   test.afterAll("Delete Test Database", async () => {
     await context?.dispose();
   });
-  }
 
   test("Add, execute, and delete user defined function", async ({ page }, testInfo) => {
     // Open container context menu and click New UDF

test/testData.ts

@@ -82,6 +82,75 @@ 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;
@@ -104,34 +173,7 @@ export async function createMultipleTestContainers({
   const creationPromises: Promise<TestContainerContext>[] = [];
 
   const databaseId = databaseName ? databaseName : generateUniqueName("db");
-  const credentials = getAzureCLICredentials();
+  const { armClient, client } = await createCosmosClientForSQLAccount(accountType);
-  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 {
@@ -158,29 +200,8 @@ 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 credentials = getAzureCLICredentials();
+  const { armClient, client } = await createCosmosClientForSQLAccount();
-  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({