96 lines
3.0 KiB
TypeScript
96 lines
3.0 KiB
TypeScript
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;
|
|
}
|
|
}
|