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;
}

export interface DocumentTestCase {
  name: string;
  databaseId: string;
  containerId: string;
  documents: TestDocument[];
}

export interface TestDocument {
  documentId: string;
  partitionKeys?: PartitionKey[];
}

export interface PartitionKey {
  key: string;
  value: string | null;
}

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 = 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;
  }
}

export const setPartitionKeys = (partitionKeys: PartitionKey[]) => {
  const result = {};

  partitionKeys.forEach((partitionKey) => {
    const { key: keyPath, value: keyValue } = partitionKey;
    const cleanPath = keyPath.startsWith("/") ? keyPath.slice(1) : keyPath;
    const keys = cleanPath.split("/");
    let current = result;

    keys.forEach((key, index) => {
      if (index === keys.length - 1) {
        current[key] = keyValue;
      } else {
        current[key] = current[key] || {};
        current = current[key];
      }
    });
  });

  return result;
};

export const serializeMongoToJson = (text: string) => {
  const normalized = text.replace(/ObjectId\("([0-9a-fA-F]{24})"\)/g, '"$1"');
  return JSON.parse(normalized);
};

export async function retry<T>(fn: () => Promise<T>, retries = 3, delayMs = 1000): Promise<T> {
  let lastError: unknown;
  for (let i = 0; i < retries; i++) {
    try {
      return await fn();
    } catch (error) {
      lastError = error;
      console.warn(`Retry ${i + 1}/${retries} failed: ${(error as Error).message}`);
      if (i < retries - 1) {
        await new Promise((res) => setTimeout(res, delayMs));
      }
    }
  }
  throw lastError;
}