Fixed unit tests

package.json

@@ -5,7 +5,7 @@
   "main": "index.js",
   "dependencies": {
     "@azure/arm-cosmosdb": "9.1.0",
-    "@azure/cosmos": "4.2.0",
+    "@azure/cosmos": "4.2.0-beta.1",
     "@azure/cosmos-language-service": "0.0.5",
     "@azure/identity": "1.5.2",
     "@azure/ms-rest-nodeauth": "3.1.1",

src/Common/IteratorUtilities.ts

@@ -1,3 +1,4 @@
+import { QueryOperationOptions } from "@azure/cosmos";
 import { QueryResults } from "../Contracts/ViewModels";
 
 interface QueryResponse {
@@ -10,13 +11,17 @@ interface QueryResponse {
 }
 
 export interface MinimalQueryIterator {
-  fetchNext: () => Promise<QueryResponse>;
+  fetchNext: (queryOperationOptions?: QueryOperationOptions) => Promise<QueryResponse>;
 }
 
 // Pick<QueryIterator<any>, "fetchNext">;
 
-export function nextPage(documentsIterator: MinimalQueryIterator, firstItemIndex: number): Promise<QueryResults> {
+export function nextPage(
-  return documentsIterator.fetchNext().then((response) => {
+  documentsIterator: MinimalQueryIterator,
+  firstItemIndex: number,
+  queryOperationOptions?: QueryOperationOptions,
+): Promise<QueryResults> {
+  return documentsIterator.fetchNext(queryOperationOptions).then((response) => {
     const documents = response.resources;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const headers = (response as any).headers || {}; // TODO this is a private key. Remove any

src/Common/dataAccess/queryDocumentsPage.ts

@@ -1,3 +1,4 @@
+import { QueryOperationOptions } from "@azure/cosmos";
 import { QueryResults } from "../../Contracts/ViewModels";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { getEntityName } from "../DocumentUtility";
@@ -8,12 +9,13 @@ export const queryDocumentsPage = async (
   resourceName: string,
   documentsIterator: MinimalQueryIterator,
   firstItemIndex: number,
+  queryOperationOptions?: QueryOperationOptions,
 ): Promise<QueryResults> => {
   const entityName = getEntityName();
   const clearMessage = logConsoleProgress(`Querying ${entityName} for container ${resourceName}`);
 
   try {
-    const result: QueryResults = await nextPage(documentsIterator, firstItemIndex);
+    const result: QueryResults = await nextPage(documentsIterator, firstItemIndex, queryOperationOptions);
     const itemCount = (result.documents && result.documents.length) || 0;
     logConsoleInfo(`Successfully fetched ${itemCount} ${entityName} for container ${resourceName}`);
     return result;

src/Common/dataAccess/updateOffer.ts

@@ -39,10 +39,7 @@ import {
   migrateTableToManualThroughput,
   updateTableThroughput,
 } from "../../Utils/arm/generatedClients/cosmos/tableResources";
-import {
+import { ThroughputSettingsUpdateParameters } from "../../Utils/arm/generatedClients/cosmos/types";
-  ThroughputSettingsGetResults,
-  ThroughputSettingsUpdateParameters,
-} from "../../Utils/arm/generatedClients/cosmos/types";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { HttpHeaders } from "../Constants";
 import { client } from "../CosmosClient";
@@ -149,28 +146,23 @@ const updateSqlContainerOffer = async (params: UpdateOfferParams): Promise<void>
   const { subscriptionId, resourceGroup, databaseAccount } = userContext;
   const accountName = databaseAccount.name;
 
-  let updatedOffer: ThroughputSettingsGetResults;
-
   if (params.migrateToAutoPilot) {
-    updatedOffer = (await migrateSqlContainerToAutoscale(
+    await migrateSqlContainerToAutoscale(
       subscriptionId,
       resourceGroup,
       accountName,
       params.databaseId,
       params.collectionId,
-    )) as ThroughputSettingsGetResults;
+    );
-    params.autopilotThroughput = updatedOffer.properties?.resource?.autoscaleSettings?.maxThroughput;
   } else if (params.migrateToManual) {
-    updatedOffer = (await migrateSqlContainerToManualThroughput(
+    await migrateSqlContainerToManualThroughput(
       subscriptionId,
       resourceGroup,
       accountName,
       params.databaseId,
       params.collectionId,
-    )) as ThroughputSettingsGetResults;
+    );
-    params.manualThroughput = updatedOffer.properties?.resource?.throughput;
+  } else {
-  }
-  if (params.throughputBuckets || !(params.migrateToAutoPilot || params.migrateToManual)) {
     const body: ThroughputSettingsUpdateParameters = createUpdateOfferBody(params);
     await updateSqlContainerThroughput(
       subscriptionId,

src/Explorer/Controls/Settings/SettingsComponent.test.tsx

@@ -1,5 +1,7 @@
+import { AuthType } from "AuthType";
 import { shallow } from "enzyme";
 import ko from "knockout";
+import { Features } from "Platform/Hosted/extractFeatures";
 import React from "react";
 import { updateCollection } from "../../../Common/dataAccess/updateCollection";
 import { updateOffer } from "../../../Common/dataAccess/updateOffer";
@@ -247,4 +249,42 @@ describe("SettingsComponent", () => {
     expect(conflictResolutionPolicy.mode).toEqual(DataModels.ConflictResolutionMode.Custom);
     expect(conflictResolutionPolicy.conflictResolutionProcedure).toEqual(expectSprocPath);
   });
+
+  it("should save throughput bucket changes when Save button is clicked", async () => {
+    updateUserContext({
+      apiType: "SQL",
+      features: { enableThroughputBuckets: true } as Features,
+      authType: AuthType.AAD,
+    });
+
+    const wrapper = shallow(<SettingsComponent {...baseProps} />);
+
+    const settingsComponentInstance = wrapper.instance() as SettingsComponent;
+    const isEnabled = settingsComponentInstance["throughputBucketsEnabled"];
+    expect(isEnabled).toBe(true);
+
+    wrapper.setState({
+      isThroughputBucketsSaveable: true,
+      throughputBuckets: [
+        { id: 1, maxThroughputPercentage: 70 },
+        { id: 2, maxThroughputPercentage: 60 },
+      ],
+    });
+
+    await settingsComponentInstance.onSaveClick();
+
+    expect(updateOffer).toHaveBeenCalledWith({
+      databaseId: collection.databaseId,
+      collectionId: collection.id(),
+      currentOffer: expect.any(Object),
+      autopilotThroughput: collection.offer().autoscaleMaxThroughput,
+      manualThroughput: collection.offer().manualThroughput,
+      throughputBuckets: [
+        { id: 1, maxThroughputPercentage: 70 },
+        { id: 2, maxThroughputPercentage: 60 },
+      ],
+    });
+
+    expect(wrapper.state("isThroughputBucketsSaveable")).toBe(false);
+  });
 });

src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.test.tsx

@@ -6,7 +6,6 @@ import { ThroughputBucketsComponent } from "./ThroughputBucketsComponent";
 describe("ThroughputBucketsComponent", () => {
   const mockOnBucketsChange = jest.fn();
   const mockOnSaveableChange = jest.fn();
-  const mockOnDiscardableChange = jest.fn();
 
   const defaultProps = {
     currentBuckets: [
@@ -19,19 +18,15 @@ describe("ThroughputBucketsComponent", () => {
     ],
     onBucketsChange: mockOnBucketsChange,
     onSaveableChange: mockOnSaveableChange,
-    onDiscardableChange: mockOnDiscardableChange,
   };
 
   beforeEach(() => {
     jest.clearAllMocks();
   });
 
-  it("renders 5 buckets with default values when input buckets are missing", () => {
+  it("renders the correct number of buckets", () => {
-    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={[{ id: 1, maxThroughputPercentage: 50 }]} />);
+    render(<ThroughputBucketsComponent {...defaultProps} />);
-
+    expect(screen.getAllByText(/Group \d+/)).toHaveLength(5);
-    expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5);
-    expect(screen.getByDisplayValue("50")).toBeInTheDocument();
-    expect(screen.getAllByDisplayValue("100").length).toBe(4);
   });
 
   it("renders buckets in the correct order even if input is unordered", () => {
@@ -41,16 +36,36 @@ describe("ThroughputBucketsComponent", () => {
     ];
     render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={unorderedBuckets} />);
 
-    const bucketLabels = screen.getAllByText(/Bucket \d+/).map((el) => el.textContent);
+    const bucketLabels = screen.getAllByText(/Group \d+/).map((el) => el.textContent);
-    expect(bucketLabels).toEqual(["Bucket 1", "Bucket 2", "Bucket 3", "Bucket 4", "Bucket 5"]);
+    expect(bucketLabels).toEqual(["Group 1 (Data Explorer Query Bucket)", "Group 2", "Group 3", "Group 4", "Group 5"]);
+  });
+
+  it("renders all provided buckets even if they exceed the max default bucket count", () => {
+    const oversizedBuckets = [
+      { id: 1, maxThroughputPercentage: 50 },
+      { id: 2, maxThroughputPercentage: 60 },
+      { id: 3, maxThroughputPercentage: 70 },
+      { id: 4, maxThroughputPercentage: 80 },
+      { id: 5, maxThroughputPercentage: 90 },
+      { id: 6, maxThroughputPercentage: 100 },
+      { id: 7, maxThroughputPercentage: 40 },
+    ];
+
+    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={oversizedBuckets} />);
+
+    expect(screen.getAllByText(/Group \d+/)).toHaveLength(7);
+
     expect(screen.getByDisplayValue("50")).toBeInTheDocument();
     expect(screen.getByDisplayValue("60")).toBeInTheDocument();
-    expect(screen.getAllByDisplayValue("100").length).toBe(3);
+    expect(screen.getByDisplayValue("70")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("80")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("90")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("100")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("40")).toBeInTheDocument();
   });
 
   it("calls onBucketsChange when a bucket value changes", () => {
     render(<ThroughputBucketsComponent {...defaultProps} />);
-
     const input = screen.getByDisplayValue("50");
     fireEvent.change(input, { target: { value: "70" } });
 
@@ -63,21 +78,12 @@ describe("ThroughputBucketsComponent", () => {
     ]);
   });
 
-  it("triggers onSaveableChange and onDiscardableChange when values change", () => {
+  it("triggers onSaveableChange when values change", () => {
     render(<ThroughputBucketsComponent {...defaultProps} />);
-
     const input = screen.getByDisplayValue("50");
     fireEvent.change(input, { target: { value: "80" } });
 
     expect(mockOnSaveableChange).toHaveBeenCalledWith(true);
-    expect(mockOnDiscardableChange).toHaveBeenCalledWith(true);
-  });
-
-  it("ensures buckets revert to default when no buckets are provided", () => {
-    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={[]} />);
-
-    expect(screen.getAllByText(/Bucket \d+/)).toHaveLength(5);
-    expect(screen.getAllByDisplayValue("100")).toHaveLength(5);
   });
 
   it("updates state consistently after multiple changes to different buckets", () => {
@@ -100,18 +106,9 @@ describe("ThroughputBucketsComponent", () => {
 
   it("resets to baseline when currentBuckets are reset", () => {
     const { rerender } = render(<ThroughputBucketsComponent {...defaultProps} />);
-
     const input1 = screen.getByDisplayValue("50");
     fireEvent.change(input1, { target: { value: "70" } });
 
-    expect(mockOnBucketsChange).toHaveBeenCalledWith([
-      { id: 1, maxThroughputPercentage: 70 },
-      { id: 2, maxThroughputPercentage: 60 },
-      { id: 3, maxThroughputPercentage: 100 },
-      { id: 4, maxThroughputPercentage: 100 },
-      { id: 5, maxThroughputPercentage: 100 },
-    ]);
-
     rerender(<ThroughputBucketsComponent {...defaultProps} currentBuckets={defaultProps.throughputBucketsBaseline} />);
 
     expect(screen.getByDisplayValue("40")).toBeInTheDocument();
@@ -120,10 +117,61 @@ describe("ThroughputBucketsComponent", () => {
 
   it("does not call onBucketsChange when value remains unchanged", () => {
     render(<ThroughputBucketsComponent {...defaultProps} />);
-
     const input = screen.getByDisplayValue("50");
     fireEvent.change(input, { target: { value: "50" } });
 
     expect(mockOnBucketsChange).not.toHaveBeenCalled();
   });
+
+  it("disables input and slider when maxThroughputPercentage is 100", () => {
+    render(
+      <ThroughputBucketsComponent
+        {...defaultProps}
+        currentBuckets={[
+          { id: 1, maxThroughputPercentage: 100 },
+          { id: 2, maxThroughputPercentage: 50 },
+        ]}
+      />,
+    );
+
+    const disabledInputs = screen.getAllByDisplayValue("100");
+    expect(disabledInputs.length).toBeGreaterThan(0);
+    expect(disabledInputs[0]).toBeDisabled();
+
+    const sliders = screen.getAllByRole("slider");
+    expect(sliders.length).toBeGreaterThan(0);
+    expect(sliders[0]).toHaveAttribute("aria-disabled", "true");
+    expect(sliders[1]).toHaveAttribute("aria-disabled", "false");
+  });
+
+  it("toggles bucket value between 50 and 100 with switch", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} />);
+    const toggles = screen.getAllByRole("switch");
+
+    fireEvent.click(toggles[0]);
+
+    expect(mockOnBucketsChange).toHaveBeenCalledWith([
+      { id: 1, maxThroughputPercentage: 100 },
+      { id: 2, maxThroughputPercentage: 60 },
+      { id: 3, maxThroughputPercentage: 100 },
+      { id: 4, maxThroughputPercentage: 100 },
+      { id: 5, maxThroughputPercentage: 100 },
+    ]);
+
+    fireEvent.click(toggles[0]);
+
+    expect(mockOnBucketsChange).toHaveBeenCalledWith([
+      { id: 1, maxThroughputPercentage: 50 },
+      { id: 2, maxThroughputPercentage: 60 },
+      { id: 3, maxThroughputPercentage: 100 },
+      { id: 4, maxThroughputPercentage: 100 },
+      { id: 5, maxThroughputPercentage: 100 },
+    ]);
+  });
+
+  it("ensures default buckets are used when no buckets are provided", () => {
+    render(<ThroughputBucketsComponent {...defaultProps} currentBuckets={[]} />);
+    expect(screen.getAllByText(/Group \d+/)).toHaveLength(5);
+    expect(screen.getAllByDisplayValue("100")).toHaveLength(5);
+  });
 });

src/Explorer/Controls/Settings/SettingsSubComponents/ThroughputInputComponents/ThroughputBucketsComponent.tsx

@@ -43,13 +43,11 @@ export const ThroughputBucketsComponent: FC<ThroughputBucketsComponentProps> = (
   useEffect(() => {
     setThroughputBuckets(getThroughputBuckets(currentBuckets));
     onSaveableChange(false);
-    // onDiscardableChange(false);
   }, [currentBuckets]);
 
   useEffect(() => {
     const isChanged = isDirty(throughputBuckets, getThroughputBuckets(throughputBucketsBaseline));
     onSaveableChange(isChanged);
-    // onDiscardableChange(isChanged);
   }, [throughputBuckets]);
 
   const handleBucketChange = (id: number, newValue: number) => {
@@ -92,10 +90,6 @@ export const ThroughputBucketsComponent: FC<ThroughputBucketsComponentProps> = (
               }}
               disabled={bucket.maxThroughputPercentage === 100}
             />
-            {/* <IconButton
-              iconProps={{ iconName: bucket.maxThroughputPercentage === 100 ? "Add" : "Remove" }}
-              onClick={() => onToggle(bucket.id, bucket.maxThroughputPercentage === 100)}
-            ></IconButton> */}
             <Toggle
               onText="Active"
               offText="Inactive"
@@ -103,12 +97,6 @@ export const ThroughputBucketsComponent: FC<ThroughputBucketsComponentProps> = (
               onChange={(event, checked) => onToggle(bucket.id, checked)}
               styles={{ root: { marginBottom: 0 }, text: { fontSize: 12 } }}
             ></Toggle>
-            {/* {bucket.id === 1 && (
-              <Stack horizontal tokens={{ childrenGap: 4 }} verticalAlign="center">
-                <Icon iconName="TagSolid" />
-                <span>Data Explorer Query Bucket</span>
-              </Stack>
-            )} */}
           </Stack>
         ))}
       </Stack>

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable no-console */
-import { FeedOptions } from "@azure/cosmos";
+import { FeedOptions, QueryOperationOptions } from "@azure/cosmos";
 import QueryError, { createMonacoErrorLocationResolver, createMonacoMarkersForQueryErrors } from "Common/QueryError";
 import { SplitterDirection } from "Common/Splitter";
 import { Platform, configContext } from "ConfigContext";
@@ -18,7 +18,7 @@ import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { KeyboardAction } from "KeyboardShortcuts";
 import { QueryConstants } from "Shared/Constants";
-import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { Allotment } from "allotment";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
@@ -368,8 +368,21 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       isExecutionError: false,
     });
 
+    let queryOperationOptions: QueryOperationOptions;
+    if (userContext.apiType === "SQL" && ruThresholdEnabled()) {
+      const ruThreshold: number = getRUThreshold();
+      queryOperationOptions = {
+        ruCapPerOperation: ruThreshold,
+      } as QueryOperationOptions;
+    }
+
     const queryDocuments = async (firstItemIndex: number) =>
-      await queryDocumentsPage(this.props.collection && this.props.collection.id(), this._iterator, firstItemIndex);
+      await queryDocumentsPage(
+        this.props.collection && this.props.collection.id(),
+        this._iterator,
+        firstItemIndex,
+        queryOperationOptions,
+      );
     this.props.tabsBaseInstance.isExecuting(true);
     this.setState({
       isExecuting: true,