Adds query comparison tests for in-memory-emulator by FabianMeiswinkel · Pull Request #4733 · Azure/azure-sdk-for-rust

FabianMeiswinkel · 2026-07-10T00:40:30Z

Summary

Adds broader Cosmos query comparison coverage for the Rust SDK in-memory emulator, including in-memory-only physical partition topology variations.

The new coverage validates that query results, resume paths, query plans, and physical partition targeting remain correct when:

hash-partitioned containers use more seed partition key values
HPK containers include more tenants, users, and sessions
the in-memory emulator runs with a single physical partition
the in-memory emulator runs with multiple physical partitions
a top-level HPK prefix is split across physical partitions, so level-2 prefixes under the same tenant route to different physical partitions

Changes

Expanded query comparison seed data:
- added more hash partition key values
- added more HPK documents under tenant-a
- added additional HPK tenants for broader fan-out coverage
Added an in-memory-only query topology test:
- query_results_match_across_physical_partition_topologies
- runs the same query scenarios across different physical partition layouts
- validates result equality for collect and resume paths
- asserts only the expected physical partitions are touched
Added physical partition targeting validation using existing wire contracts:
- uses raw /pkranges responses to determine expected physical ranges
- uses query page x-ms-documentdb-partitionkeyrangeid response headers to determine actual touched ranges
- avoids adding emulator-side touch recording
Added a test-only in-memory emulator split hook:
- EmulatorStore::split_partition_at_epk
- allows tests to create deterministic split points, including inside a top-level HPK prefix
- existing split_partition behavior remains midpoint-based and unchanged for existing callers
Preserved live/external comparison behavior:
- external result drains remain skipped for partial HPK prefix scenarios where the standard gateway rejects partial HPK EPK-range execution
- query plan comparison still runs for those scenarios

Validation

Live account comparison passed:

$oldMode = $env:AZURE_COSMOS_TEST_MODE
$env:AZURE_COSMOS_TEST_MODE = 'required'
cargo test -p azure_data_cosmos --test in_memory_emulator --all-features query_results_plans_and_resume_paths_match -- --nocapture

Local Cosmos emulator comparison passed against https://127.0.0.1:8888:

$oldMode = $env:AZURE_COSMOS_TEST_MODE
$oldConnection = $env:AZURE_COSMOS_CONNECTION_STRING
$env:AZURE_COSMOS_TEST_MODE = 'required'
$env:AZURE_COSMOS_CONNECTION_STRING = 'AccountEndpoint=https://127.0.0.1:8888;AccountKey=<redacted>;'
cargo test -p azure_data_cosmos --test in_memory_emulator --all-features query_results_plans_and_resume_paths_match -- --nocapture

In-memory-only custom split topology coverage passed:

$oldMode = $env:AZURE_COSMOS_TEST_MODE
$env:AZURE_COSMOS_TEST_MODE = 'skipped'
cargo test -p azure_data_cosmos --test in_memory_emulator --all-features query_results_match_across_physical_partition_topologies -- --nocapture

Focused driver query-plan and split coverage passed:

cargo test -p azure_data_cosmos_driver --test in_memory_emulator --all-features query_plan -- --nocapture
cargo test -p azure_data_cosmos_driver --test in_memory_emulator --all-features split -- --nocapture

Clippy passed for the touched test targets:

cargo clippy -p azure_data_cosmos --test in_memory_emulator --all-features
cargo clippy -p azure_data_cosmos_driver --test in_memory_emulator --all-features

Copilot

Pull request overview

Adds query parity and topology tests for the Cosmos in-memory emulator.

Changes:

Adds query result, plan, continuation, and partition-topology comparisons.
Improves HPK prefix routing and query response metadata.
Adds explicit EPK partition splitting for tests.

Show a summary per file

File	Description
`azure_data_cosmos/tests/in_memory_emulator_tests/validation.rs`	Adds query header validation.
`azure_data_cosmos/tests/in_memory_emulator_tests/query_comparison.rs`	Adds query comparison scenarios.
`azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs`	Registers the new tests.
`azure_data_cosmos_driver/tests/in_memory_emulator_tests/query.rs`	Expands query-plan assertions.
`azure_data_cosmos_driver/src/query/plan/mod.rs`	Supports leading HPK-prefix extraction.
`azure_data_cosmos_driver/src/in_memory_emulator/store.rs`	Supports explicit EPK splits.
`azure_data_cosmos_driver/src/in_memory_emulator/operations.rs`	Improves query plans and response headers.
`azure_data_cosmos_driver/src/driver/cosmos_driver.rs`	Exposes cached regions to comparison tests.