Assert explain hash fields through assertProperties

documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_pipeline.py

@@ -20,54 +20,34 @@
 
 import pytest
 
-from documentdb_tests.framework.assertions import assertProperties
+from documentdb_tests.compatibility.tests.system.diagnostic.utils.explain_hash import (
+    PLAN_CACHE_SHAPE_HASH,
+    QUERY_SHAPE_HASH,
+    assert_hash_relation,
+)
 from documentdb_tests.framework.executor import execute_command
 from documentdb_tests.framework.parametrize import pytest_params
-from documentdb_tests.framework.property_checks import Eq, Ne
 from documentdb_tests.framework.test_case import BaseTestCase
 
 pytestmark = pytest.mark.admin
 
 
-def _get_plan_cache_shape_hash(collection, pipeline: list) -> str:
-    """Return queryPlanner.planCacheShapeHash from explain for an aggregate pipeline."""
-    cmd: dict[str, Any] = {
-        "aggregate": collection.name,
-        "pipeline": pipeline,
-        "cursor": {},
-    }
-    result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"})
-    if isinstance(result, Exception):
-        raise RuntimeError(f"explain command failed: {result}")
-    # Standard format: queryPlanner at top level.
-    query_planner = result.get("queryPlanner", {})
-    if "planCacheShapeHash" in query_planner:
-        return str(query_planner["planCacheShapeHash"])
-    # Stages format ($out/$merge): nested under stages.0.$cursor.queryPlanner.
-    stages = result.get("stages", [])
-    if stages:
-        cursor_stage = stages[0].get("$cursor", {})
-        qp = cursor_stage.get("queryPlanner", {})
-        if "planCacheShapeHash" in qp:
-            return str(qp["planCacheShapeHash"])
-    raise RuntimeError(
-        f"explain did not return planCacheShapeHash in expected location; got: {result}"
-    )
+# Explain reports planCacheShapeHash under queryPlanner for most pipelines, and
+# under stages.0.$cursor.queryPlanner when the pipeline ends in $out or $merge.
+PIPELINE_PLAN_CACHE_SHAPE_HASH = (
+    PLAN_CACHE_SHAPE_HASH,
+    f"stages.0.$cursor.{PLAN_CACHE_SHAPE_HASH}",
+)
 
 
-def _get_query_shape_hash(collection, pipeline: list) -> str:
-    """Return the top-level queryShapeHash from explain for an aggregate pipeline."""
+def _explain_aggregate(collection, pipeline: list) -> Any:
+    """Run explain on an aggregate pipeline and return the raw result (or the exception)."""
     cmd: dict[str, Any] = {
         "aggregate": collection.name,
         "pipeline": pipeline,
         "cursor": {},
     }
-    result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"})
-    if isinstance(result, Exception):
-        raise RuntimeError(f"explain command failed: {result}")
-    if "queryShapeHash" not in result:
-        raise RuntimeError(f"explain did not return queryShapeHash; got: {result}")
-    return str(result["queryShapeHash"])
+    return execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"})
 
 
 @dataclass(frozen=True)
@@ -361,14 +341,10 @@ def test_pipeline_plancacheshapehash_stage_sensitivity(collection, test):
     collection.insert_many(
         [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "sub": {"x": i}} for i in range(20)]
     )
-    hash_a = _get_plan_cache_shape_hash(collection, list(test.pipeline_a))
-    hash_b = _get_plan_cache_shape_hash(collection, list(test.pipeline_b))
-    check = Eq(hash_b) if test.plan_cache_same else Ne(hash_b)
-    assertProperties(
-        {"planCacheShapeHash": hash_a},
-        {"planCacheShapeHash": check},
-        msg=test.msg,
-        raw_res=True,
+    result_a = _explain_aggregate(collection, list(test.pipeline_a))
+    result_b = _explain_aggregate(collection, list(test.pipeline_b))
+    assert_hash_relation(
+        result_a, result_b, PIPELINE_PLAN_CACHE_SHAPE_HASH, same=test.plan_cache_same, msg=test.msg
     )
 
 
@@ -378,12 +354,8 @@ def test_pipeline_queryshapehash_considers_all_stages(collection, test):
     collection.insert_many(
         [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "sub": {"x": i}} for i in range(20)]
     )
-    hash_a = _get_query_shape_hash(collection, list(test.pipeline_a))
-    hash_b = _get_query_shape_hash(collection, list(test.pipeline_b))
-    check = Eq(hash_b) if test.query_shape_same else Ne(hash_b)
-    assertProperties(
-        {"queryShapeHash": hash_a},
-        {"queryShapeHash": check},
-        msg=test.msg,
-        raw_res=True,
+    result_a = _explain_aggregate(collection, list(test.pipeline_a))
+    result_b = _explain_aggregate(collection, list(test.pipeline_b))
+    assert_hash_relation(
+        result_a, result_b, QUERY_SHAPE_HASH, same=test.query_shape_same, msg=test.msg
     )

documentdb_tests/compatibility/tests/system/diagnostic/commands/explain/test_explain_queryhash_shape.py

@@ -10,60 +10,34 @@
 
 import pytest
 
-from documentdb_tests.framework.assertions import assertProperties
+from documentdb_tests.compatibility.tests.system.diagnostic.utils.explain_hash import (
+    PLAN_CACHE_SHAPE_HASH,
+    QUERY_SHAPE_HASH,
+    assert_hash_relation,
+)
 from documentdb_tests.framework.executor import execute_command
 from documentdb_tests.framework.parametrize import pytest_params
-from documentdb_tests.framework.property_checks import Eq, Ne
 from documentdb_tests.framework.test_case import BaseTestCase
 
 pytestmark = pytest.mark.admin
 
 
-def get_query_hash(
+def explain_find(
     collection,
     find_filter: dict,
     sort: Optional[dict] = None,
     projection: Optional[dict] = None,
     collation: Optional[dict] = None,
-) -> str:
-    """Return queryPlanner.planCacheShapeHash from explain for a find filter."""
+) -> Any:
+    """Run explain on a find and return the raw result (or the exception)."""
     cmd: dict[str, Any] = {"find": collection.name, "filter": find_filter}
     if sort is not None:
         cmd["sort"] = sort
     if projection is not None:
         cmd["projection"] = projection
     if collation is not None:
         cmd["collation"] = collation
-    result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"})
-    if isinstance(result, Exception):
-        raise RuntimeError(f"explain command failed: {result}")
-    query_planner = result.get("queryPlanner", {})
-    if "planCacheShapeHash" not in query_planner:
-        raise RuntimeError(f"explain did not return queryPlanner.planCacheShapeHash; got: {result}")
-    return str(query_planner["planCacheShapeHash"])
-
-
-def get_query_shape_hash(
-    collection,
-    find_filter: dict,
-    sort: Optional[dict] = None,
-    projection: Optional[dict] = None,
-    collation: Optional[dict] = None,
-) -> str:
-    """Return the top-level queryShapeHash from explain for a find filter."""
-    cmd: dict[str, Any] = {"find": collection.name, "filter": find_filter}
-    if sort is not None:
-        cmd["sort"] = sort
-    if projection is not None:
-        cmd["projection"] = projection
-    if collation is not None:
-        cmd["collation"] = collation
-    result = execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"})
-    if isinstance(result, Exception):
-        raise RuntimeError(f"explain command failed: {result}")
-    if "queryShapeHash" not in result:
-        raise RuntimeError(f"explain did not return queryShapeHash; got: {result}")
-    return str(result["queryShapeHash"])
+    return execute_command(collection, {"explain": cmd, "verbosity": "queryPlanner"})
 
 
 @dataclass(frozen=True)
@@ -297,18 +271,13 @@ def test_explain_queryhash_same_shape(collection, test):
     collection.insert_many(
         [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)]
     )
-    hash_a = get_query_hash(
+    result_a = explain_find(
         collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a
     )
-    hash_b = get_query_hash(
+    result_b = explain_find(
         collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b
     )
-    assertProperties(
-        {"planCacheShapeHash": hash_a},
-        {"planCacheShapeHash": Eq(hash_b)},
-        msg=test.msg,
-        raw_res=True,
-    )
+    assert_hash_relation(result_a, result_b, PLAN_CACHE_SHAPE_HASH, same=True, msg=test.msg)
 
 
 DIFF_SHAPE_TESTS: list[QueryHashCase] = [
@@ -516,18 +485,13 @@ def test_explain_queryhash_different_shape(collection, test):
     collection.insert_many(
         [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)]
     )
-    hash_a = get_query_hash(
+    result_a = explain_find(
         collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a
     )
-    hash_b = get_query_hash(
+    result_b = explain_find(
         collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b
     )
-    assertProperties(
-        {"planCacheShapeHash": hash_a},
-        {"planCacheShapeHash": Ne(hash_b)},
-        msg=test.msg,
-        raw_res=True,
-    )
+    assert_hash_relation(result_a, result_b, PLAN_CACHE_SHAPE_HASH, same=False, msg=test.msg)
 
 
 @pytest.mark.parametrize("test", pytest_params(SAME_SHAPE_TESTS))
@@ -542,16 +506,13 @@ def test_explain_queryshapehash_same_shape(collection, test):
     collection.insert_many(
         [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)]
     )
-    hash_a = get_query_shape_hash(
+    result_a = explain_find(
         collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a
     )
-    hash_b = get_query_shape_hash(
+    result_b = explain_find(
         collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b
     )
-    check = Eq(hash_b) if expected_same else Ne(hash_b)
-    assertProperties(
-        {"queryShapeHash": hash_a}, {"queryShapeHash": check}, msg=test.msg, raw_res=True
-    )
+    assert_hash_relation(result_a, result_b, QUERY_SHAPE_HASH, same=expected_same, msg=test.msg)
 
 
 @pytest.mark.parametrize("test", pytest_params(DIFF_SHAPE_TESTS))
@@ -565,13 +526,10 @@ def test_explain_queryshapehash_different_shape(collection, test):
     collection.insert_many(
         [{"_id": i, "a": i % 5, "b": i % 3, "arr": [i, i + 1], "s": str(i)} for i in range(20)]
     )
-    hash_a = get_query_shape_hash(
+    result_a = explain_find(
         collection, test.filter_a, test.sort_a, test.projection_a, test.collation_a
     )
-    hash_b = get_query_shape_hash(
+    result_b = explain_find(
         collection, test.filter_b, test.sort_b, test.projection_b, test.collation_b
     )
-    check = Eq(hash_b) if expected_same else Ne(hash_b)
-    assertProperties(
-        {"queryShapeHash": hash_a}, {"queryShapeHash": check}, msg=test.msg, raw_res=True
-    )
+    assert_hash_relation(result_a, result_b, QUERY_SHAPE_HASH, same=expected_same, msg=test.msg)

documentdb_tests/compatibility/tests/system/diagnostic/utils/explain_hash.py

@@ -0,0 +1,78 @@
+"""Helpers for comparing the query hashes that explain reports.
+
+The hash tests run explain on two queries and assert that their hashes match
+(same shape) or differ (different shape). The helpers here keep the explain
+result raw and do every check, including that the hash is present at all,
+through a single ``assertProperties`` call. A failed explain or a missing hash
+field is therefore reported as an assertion failure instead of an error raised
+before the assertion.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Union
+
+from documentdb_tests.framework.assertions import assertProperties
+from documentdb_tests.framework.property_checks import Eq, Ne, NonEmptyStr
+
+PLAN_CACHE_SHAPE_HASH = "queryPlanner.planCacheShapeHash"
+"""Dotted path of the plan cache shape hash in a standard explain result."""
+
+QUERY_SHAPE_HASH = "queryShapeHash"
+"""Dotted path of the top-level query shape hash in an explain result."""
+
+
+def _lookup(result: Any, path: str) -> Any:
+    """Return the value at a dotted path (dict keys and list indices), or None when absent."""
+    value: Any = result
+    for part in path.split("."):
+        if isinstance(value, list) and part.isdigit() and int(part) < len(value):
+            value = value[int(part)]
+        elif isinstance(value, dict) and part in value:
+            value = value[part]
+        else:
+            return None
+    return value
+
+
+def assert_hash_relation(
+    result_a: Any,
+    result_b: Any,
+    paths: Union[str, tuple[str, ...]],
+    same: bool,
+    msg: str,
+) -> None:
+    """Assert both explain results report a hash, and that the two hashes match or differ.
+
+    Args:
+        result_a: Raw explain result (or exception) for the first query.
+        result_b: Raw explain result (or exception) for the second query.
+        paths: Dotted path of the hash, or several candidate paths tried in
+            order (for explain output whose layout depends on the pipeline).
+        same: True to assert the hashes are equal, False to assert they differ.
+        msg: Failure message.
+    """
+    candidates = (paths,) if isinstance(paths, str) else paths
+    for result in (result_a, result_b):
+        if isinstance(result, Exception):
+            # The failed explain is the result under test: assert on it and stop,
+            # so it is reported once rather than again as two missing hashes.
+            assertProperties(result, {candidates[0]: NonEmptyStr()}, msg=msg, raw_res=True)
+            return
+    hashes: dict[str, str] = {}
+    for key, result in (("hash_a", result_a), ("hash_b", result_b)):
+        for path in candidates:
+            value = _lookup(result, path)
+            if value is not None:
+                hashes[key] = str(value)
+                break
+    # Compare only when the second hash exists; its absence is reported on its own.
+    checks_a: list[Any] = [NonEmptyStr()]
+    if "hash_b" in hashes:
+        checks_a.append(Eq(hashes["hash_b"]) if same else Ne(hashes["hash_b"]))
+    assertProperties(
+        hashes,
+        {"hash_a": checks_a, "hash_b": NonEmptyStr()},
+        msg=f"{msg} (hash at {' or '.join(candidates)})",
+        raw_res=True,
+    )