Updated Test.

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml

@@ -1,6 +1,39 @@
 <Results>
  <Result>
- <Input />
+ <Input>
+ <Description>SELECT</Description>
+ <Query><![CDATA[SELECT c.id, c.value2 FROM c]]></Query>
+ <ODE>True</ODE>
+ </Input>
+ <Output>
+ <Documents><![CDATA[{"id":"2","value2":"12"},
+{"id":"2","value2":"17"},
+{"id":"2","value2":"2"},
+{"id":"2","value2":"22"},
+{"id":"2","value2":"27"},
+{"id":"2","value2":"32"},
+{"id":"2","value2":"37"},
+{"id":"2","value2":"42"},
+{"id":"2","value2":"47"},
+{"id":"2","value2":"52"},
+{"id":"2","value2":"57"},
+{"id":"2","value2":"62"},
+{"id":"2","value2":"67"},
+{"id":"2","value2":"7"},
+{"id":"2","value2":"72"},
+{"id":"2","value2":"77"},
+{"id":"2","value2":"82"},
+{"id":"2","value2":"87"},
+{"id":"2","value2":"92"},
+{"id":"2","value2":"97"}]]></Documents>
+ </Output>
+ </Result>
+ <Result>
+ <Input>
+ <Description>SELECT without ODE</Description>
+ <Query><![CDATA[SELECT c.id, c.value2 FROM c]]></Query>
+ <ODE>False</ODE>
+ </Input>
  <Output>
  <Documents><![CDATA[{"id":"2","value2":"12"},
 {"id":"2","value2":"17"},

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs

@@ -33,7 +33,12 @@ public class SubpartitionTests : BaselineTests<SubpartitionTestInput, Subpartiti
  [TestMethod]
  public void TestQueriesOnSplitContainer()
  {
- this.ExecuteTestSuite(new List<SubpartitionTestInput> { new SubpartitionTestInput("Test Queries on Split container") });
+ List<SubpartitionTestInput> inputs = new List<SubpartitionTestInput>
+ {
+ new SubpartitionTestInput("SELECT", query: @"SELECT c.id, c.value2 FROM c", ode: true),
+ new SubpartitionTestInput("SELECT without ODE", query: @"SELECT c.id, c.value2 FROM c", ode: false),
+ };
+ this.ExecuteTestSuite(inputs);
  }
 
  /// <summary>
@@ -58,11 +63,6 @@ public async Task VerifyTestFrameworkSupportsPartitionSplit()
  Assert.IsTrue(pkRangeId.Exception.InnerException.Message.StartsWith("Epk Range: [B5-D7-B7-26-D6-EA-DB-11-F1-EF-AD-92-12-15-D6-60,B5-D7-B7-26-D6-EA-DB-11-F1-EF-AD-92-12-15-D6-60-FF) is gone."), "Gone exception is expected!");
  }
 
- public SubpartitionTestOutput ExecuteTest2(SubpartitionTestInput input)
- {
- return new SubpartitionTestOutput(new List<CosmosElement>());
- }
-
  public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
  {
  IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount).Result;
@@ -74,7 +74,7 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
  PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build()
  };
  (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) =
- CreateInputParamsAndQueryContext(queryRequestOptions);
+ CreateInputParamsAndQueryContext(input, queryRequestOptions);
  IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create(
  documentContainer,
  cosmosQueryContextCore,
@@ -95,9 +95,9 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
  return new SubpartitionTestOutput(documents);
  }
 
- private static Tuple<CosmosQueryExecutionContextFactory.InputParameters, CosmosQueryContextCore> CreateInputParamsAndQueryContext(QueryRequestOptions queryRequestOptions, bool clientDisableOde = false)
+ private static Tuple<CosmosQueryExecutionContextFactory.InputParameters, CosmosQueryContextCore> CreateInputParamsAndQueryContext(SubpartitionTestInput input, QueryRequestOptions queryRequestOptions)
  {
- string query = @"SELECT c.id, c.value2 FROM c";
+ string query = input.Query;
  CosmosElement continuationToken = null;
  PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition()
  {
@@ -111,11 +111,13 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
  Version = PartitionKeyDefinitionVersion.V2,
  };
 
+ queryRequestOptions.EnableOptimisticDirectExecution = input.ODE;
+
  CosmosSerializerCore serializerCore = new();
  using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(new SqlQuerySpec(query), Documents.ResourceType.Document));
  string sqlQuerySpecJsonString = streamReader.ReadToEnd();
 
- (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition, clientDisableOde);
+ (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition);
  CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters(
  sqlQuerySpec: new SqlQuerySpec(query),
  initialUserContinuationToken: continuationToken,
@@ -147,9 +149,9 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
  return Tuple.Create(inputParameters, cosmosQueryContextCore);
  }
 
- internal static Tuple<PartitionedQueryExecutionInfo, QueryPartitionProvider> GetPartitionedQueryExecutionInfoAndPartitionProvider(string querySpecJsonString, PartitionKeyDefinition pkDefinition, bool clientDisableOde = false)
+ internal static Tuple<PartitionedQueryExecutionInfo, QueryPartitionProvider> GetPartitionedQueryExecutionInfoAndPartitionProvider(string querySpecJsonString, PartitionKeyDefinition pkDefinition)
  {
- QueryPartitionProvider queryPartitionProvider = CreateCustomQueryPartitionProvider("clientDisableOptimisticDirectExecution", clientDisableOde.ToString().ToLower());
+ QueryPartitionProvider queryPartitionProvider = CreateCustomQueryPartitionProvider();
  TryCatch<PartitionedQueryExecutionInfo> tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo(
  querySpecJsonString: querySpecJsonString,
  partitionKeyDefinition: pkDefinition,
@@ -165,7 +167,7 @@ internal static Tuple<PartitionedQueryExecutionInfo, QueryPartitionProvider> Get
  return Tuple.Create(partitionedQueryExecutionInfo, queryPartitionProvider);
  }
 
- private static QueryPartitionProvider CreateCustomQueryPartitionProvider(string key, string value)
+ private static QueryPartitionProvider CreateCustomQueryPartitionProvider()
  {
  Dictionary<string, object> queryEngineConfiguration = new Dictionary<string, object>()
  {
@@ -192,8 +194,6 @@ private static QueryPartitionProvider CreateCustomQueryPartitionProvider(string
  {"clientDisableOptimisticDirectExecution", false}
  };
 
- queryEngineConfiguration[key] = bool.TryParse(value, out bool boolValue) ? boolValue : value;
-
  return new QueryPartitionProvider(queryEngineConfiguration);
  }
 
@@ -227,7 +227,6 @@ private static async Task<InMemoryContainer> CreateSplitInMemoryDocumentContaine
  InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true);
  for (int i = 0; i < numItems; i++)
  {
- // Insert an item
  CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\" }}");
  while (true)
  {
@@ -351,13 +350,24 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
 
  public class SubpartitionTestInput : BaselineTestInput
  {
- public SubpartitionTestInput(string description)
+ public SubpartitionTestInput(string description, string query, bool ode)
  :base(description)
  {
+ this.Query = query;
+ this.ODE = ode;
  }
 
+ internal string Query { get; }
+
+ internal bool ODE { get; }
+
  public override void SerializeAsXml(XmlWriter xmlWriter)
  {
+ xmlWriter.WriteElementString("Description", this.Description);
+ xmlWriter.WriteStartElement("Query");
+ xmlWriter.WriteCData(this.Query);
+ xmlWriter.WriteEndElement();
+ xmlWriter.WriteElementString("ODE", this.ODE.ToString());
  }
  }