update test

Azure · Jan 13, 2025 · dc7bfff · dc7bfff
1 parent 4dcbb35
commit dc7bfff
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 49 deletions.
diff --git a/sdk/cosmosdb/cosmos/review/cosmos.api.md b/sdk/cosmosdb/cosmos/review/cosmos.api.md
@@ -1054,6 +1054,7 @@ export interface FeedOptions extends SharedOptions {
     continuationToken?: string;
     continuationTokenLimitInKB?: number;
     disableNonStreamingOrderByQuery?: boolean;
+    enableQueryControl?: boolean;
     enableScanInQuery?: boolean;
     forceQueryPlan?: boolean;
     maxDegreeOfParallelism?: number;

diff --git a/sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts b/sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts
@@ -429,7 +429,7 @@ export abstract class ParallelQueryExecutionContextBase implements ExecutionCont
             documentProducer: DocumentProducer,
           ): Promise<void> => {
             try {
-              await documentProducer.bufferMore(this.getDiagnosticNode());
+              await documentProducer.bufferMore(diagnosticNode);
               // if buffer of document producer is filled, add it to the buffered document producers queue
               const nextItem = documentProducer.peakNextItem();
               if (nextItem !== undefined) {
@@ -443,10 +443,7 @@ export abstract class ParallelQueryExecutionContextBase implements ExecutionCont
                 // So that later parts of the code can repair the execution context
                 // refresh the partition key ranges and ctreate new document producers and add it to the queue
 
-                await this._enqueueReplacementDocumentProducers(
-                  this.getDiagnosticNode(),
-                  documentProducer,
-                );
+                await this._enqueueReplacementDocumentProducers(diagnosticNode, documentProducer);
                 resolve();
               } else {
                 this.err = err;

diff --git a/sdk/cosmosdb/cosmos/test/internal/unit/common/TestOrderbyQueryExecutionContext.ts b/sdk/cosmosdb/cosmos/test/internal/unit/common/TestOrderbyQueryExecutionContext.ts
diff --git a/sdk/cosmosdb/cosmos/test/internal/unit/query/pipelinedQueryExecutionContext.spec.ts b/sdk/cosmosdb/cosmos/test/internal/unit/query/pipelinedQueryExecutionContext.spec.ts
@@ -84,7 +84,7 @@ describe("PipelineQueryExecutionContext", function () {
       );
       // Mock the endpoint's fetchMore method to return 3 documents in every call
       context["endpoint"] = {
-        fetchMore: async (diagnosticNode: any) => {
+        fetchMore: async () => {
           return {
             result: [
               createMockDocument("1", "doc1", "value1"),
@@ -123,7 +123,7 @@ describe("PipelineQueryExecutionContext", function () {
       );
       let i = 0;
       context["endpoint"] = {
-        fetchMore: async (diagnosticNode: any) => {
+        fetchMore: async () => {
           if (i < 3) {
             i++;
             return {
@@ -166,7 +166,7 @@ describe("PipelineQueryExecutionContext", function () {
       );
       let i = 0;
       context["endpoint"] = {
-        fetchMore: async (diagnosticNode: any) => {
+        fetchMore: async () => {
           if (i < 1) {
             i++;
             return {
@@ -207,7 +207,7 @@ describe("PipelineQueryExecutionContext", function () {
         false,
       );
       context["endpoint"] = {
-        fetchMore: async (diagnosticNode: any) => {
+        fetchMore: async () => {
           return {
             result: undefined,
             headers: {},
@@ -237,7 +237,7 @@ describe("PipelineQueryExecutionContext", function () {
       );
       let i = 0;
       context["endpoint"] = {
-        fetchMore: async (diagnosticNode: any) => {
+        fetchMore: async () => {
           if (i < 1) {
             i++;
             return {
@@ -282,7 +282,7 @@ describe("PipelineQueryExecutionContext", function () {
         true,
       );
       context["endpoint"] = {
-        fetchMore: async (diagnosticNode: any) => {
+        fetchMore: async () => {
           return {
             result: [
               createMockDocument("1", "doc1", "value1"),

diff --git a/sdk/cosmosdb/cosmos/test/public/functional/test-query.spec.ts b/sdk/cosmosdb/cosmos/test/public/functional/test-query.spec.ts