cosmos-explorer/test/testData.ts

import { CosmosDBManagementClient } from "@azure/arm-cosmosdb";
import { BulkOperationType, Container, CosmosClient, Database, JSONObject } from "@azure/cosmos";
import crypto from "crypto";
import {
  TestAccount,
  generateUniqueName,
  getAccountName,
  getAzureCLICredentials,
  resourceGroupName,
  subscriptionId,
} from "./fx";

export interface TestItem {
  id: string;
  partitionKey: string;
  randomData: string;
}

const partitionCount = 4;

// If we increase this number, we need to split bulk creates into multiple batches.
// Bulk operations are limited to 100 items per partition.
const itemsPerPartition = 100;

function createTestItems(): TestItem[] {
  const items: TestItem[] = [];
  for (let i = 0; i < partitionCount; i++) {
    for (let j = 0; j < itemsPerPartition; j++) {
      const id = crypto.randomBytes(32).toString("base64");
      items.push({
        id,
        partitionKey: `partition_${i}`,
        randomData: crypto.randomBytes(32).toString("base64"),
      });
    }
  }
  return items;
}

export const TestData: TestItem[] = createTestItems();

export class TestContainerContext {
  constructor(
    public armClient: CosmosDBManagementClient,
    public client: CosmosClient,
    public database: Database,
    public container: Container,
    public testData: Map<string, TestItem>,
  ) {}

  async dispose() {
    await this.database.delete();
  }
}

export async function createTestSQLContainer(includeTestData?: boolean) {
  const databaseId = generateUniqueName("db");
  const containerId = "testcontainer"; // A unique container name isn't needed because the database is unique
  const credentials = await getAzureCLICredentials();
  const armClient = new CosmosDBManagementClient(credentials, subscriptionId);
  const accountName = getAccountName(TestAccount.SQL);
  const account = await armClient.databaseAccounts.get(resourceGroupName, accountName);
  const keys = await armClient.databaseAccounts.listKeys(resourceGroupName, accountName);
  const client = new CosmosClient({
    endpoint: account.documentEndpoint!,
    key: keys.primaryMasterKey,
  });
  const { database } = await client.databases.createIfNotExists({ id: databaseId });
  try {
    const { container } = await database.containers.createIfNotExists({
      id: containerId,
      partitionKey: "/partitionKey",
    });
    if (includeTestData) {
      const batchCount = TestData.length / 100;
      for (let i = 0; i < batchCount; i++) {
        const batchItems = TestData.slice(i * 100, i * 100 + 100);
        await container.items.bulk(
          batchItems.map((item) => ({
            operationType: BulkOperationType.Create,
            resourceBody: item as unknown as JSONObject,
          })),
        );
      }
    }

    const testDataMap = new Map<string, TestItem>();
    TestData.forEach((item) => testDataMap.set(item.id, item));

    return new TestContainerContext(armClient, client, database, container, testDataMap);
  } catch (e) {
    await database.delete();
    throw e;
  }
}