From abd41485e899b384ba8c6541c01e488643c20773 Mon Sep 17 00:00:00 2001 From: Sindhu Balasubramanian Date: Wed, 14 Jan 2026 22:22:17 -0800 Subject: [PATCH] Address comments --- .../scaleAndSettings/sharedThroughput.spec.ts | 66 +--------------- test/testData.ts | 79 ++++++------------- 2 files changed, 24 insertions(+), 121 deletions(-) diff --git a/test/sql/scaleAndSettings/sharedThroughput.spec.ts b/test/sql/scaleAndSettings/sharedThroughput.spec.ts index bbe1adb89..64b1dc09d 100644 --- a/test/sql/scaleAndSettings/sharedThroughput.spec.ts +++ b/test/sql/scaleAndSettings/sharedThroughput.spec.ts @@ -1,6 +1,3 @@ -import { CosmosDBManagementClient } from "@azure/arm-cosmosdb"; -import { CosmosClient, CosmosClientOptions, Database } from "@azure/cosmos"; -import { AzureIdentityCredentialAdapter } from "@azure/ms-rest-js"; import { Locator, expect, test } from "@playwright/test"; import { CommandBarButton, @@ -9,69 +6,8 @@ import { TEST_AUTOSCALE_MAX_THROUGHPUT_RU_4K, TEST_MANUAL_THROUGHPUT_RU, TestAccount, - generateUniqueName, - getAccountName, - getAzureCLICredentials, - resourceGroupName, - subscriptionId, } from "../../fx"; - -// Helper class for database context -class TestDatabaseContext { - constructor( - public armClient: CosmosDBManagementClient, - public client: CosmosClient, - public database: Database, - ) {} - - async dispose() { - await this.database.delete(); - } -} - -// Options for creating test database -interface CreateTestDBOptions { - throughput?: number; - maxThroughput?: number; // For autoscale -} - -// Helper function to create a test database with shared throughput -async function createTestDB(options?: CreateTestDBOptions): Promise { - const databaseId = generateUniqueName("db"); - 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 => { - 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); - - // 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); -} +import { TestDatabaseContext, createTestDB } from "../../testData"; test.describe("Database with Shared Throughput", () => { let dbContext: TestDatabaseContext = null!; diff --git a/test/testData.ts b/test/testData.ts index fed9b0703..6d892cc60 100644 --- a/test/testData.ts +++ b/test/testData.ts @@ -99,23 +99,31 @@ export interface CreateTestDBOptions { maxThroughput?: number; // For autoscale } -export async function createTestDB(options?: CreateTestDBOptions): Promise { - const databaseId = generateUniqueName("db"); +// 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(TestAccount.SQL); + const accountName = getAccountName(accountType); const account = await armClient.databaseAccounts.get(resourceGroupName, accountName); const clientOptions: CosmosClientOptions = { endpoint: account.documentEndpoint!, }; - const nosqlAccountRbacToken = process.env.NOSQL_TESTACCOUNT_TOKEN; - if (nosqlAccountRbacToken) { + 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 => { const AUTH_PREFIX = `type=aad&ver=1.0&sig=`; - const authorizationToken = `${AUTH_PREFIX}${nosqlAccountRbacToken}`; + const authorizationToken = `${AUTH_PREFIX}${rbacToken}`; return authorizationToken; }; } else { @@ -125,6 +133,13 @@ export async function createTestDB(options?: CreateTestDBOptions): Promise { + 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({ @@ -158,34 +173,7 @@ export async function createMultipleTestContainers({ const creationPromises: Promise[] = []; const databaseId = databaseName ? databaseName : generateUniqueName("db"); - 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 => { - 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 { armClient, client } = await createCosmosClientForSQLAccount(accountType); const { database } = await client.databases.createIfNotExists({ id: databaseId }); try { @@ -212,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 adaptedCredentials = new AzureIdentityCredentialAdapter(credentials); - const armClient = new CosmosDBManagementClient(adaptedCredentials, subscriptionId); - const accountName = getAccountName(TestAccount.SQL); - const account = await armClient.databaseAccounts.get(resourceGroupName, accountName); + const { armClient, client } = await createCosmosClientForSQLAccount(); - const clientOptions: CosmosClientOptions = { - endpoint: account.documentEndpoint!, - }; - - const nosqlAccountRbacToken = process.env.NOSQL_TESTACCOUNT_TOKEN; - if (nosqlAccountRbacToken) { - clientOptions.tokenProvider = async (): Promise => { - 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({