Report setup and intermediate command failures through the assertion

documentdb_tests/compatibility/tests/core/indexes/commands/listIndexes/test_listIndexes_cursor_batching.py

@@ -96,15 +96,21 @@ def test_listIndexes_getMore_retrieves_remaining(collection_with_3_indexes):
         collection_with_3_indexes,
         {"listIndexes": collection_with_3_indexes.name, "cursor": {"batchSize": 1}},
     )
-    cursor_id = result["cursor"]["id"]
-
-    get_more_result = execute_command(
-        collection_with_3_indexes,
-        {"getMore": cursor_id, "collection": collection_with_3_indexes.name},
-    )
+    if isinstance(result, Exception):
+        outcome = result
+    else:
+        get_more_result = execute_command(
+            collection_with_3_indexes,
+            {"getMore": result["cursor"]["id"], "collection": collection_with_3_indexes.name},
+        )
+        outcome = (
+            get_more_result
+            if isinstance(get_more_result, Exception)
+            else {"cursor": {"firstBatch": get_more_result["cursor"]["nextBatch"]}, "ok": 1.0}
+        )
 
     assertSuccess(
-        {"cursor": {"firstBatch": get_more_result["cursor"]["nextBatch"]}, "ok": 1.0},
+        outcome,
         [{"v": 2, "key": {"a": 1}, "name": "a_1"}, {"v": 2, "key": {"b": 1}, "name": "b_1"}],
     )
 
@@ -141,19 +147,33 @@ def test_listIndexes_batchSize_1_iterate_all(collection_with_3_indexes):
         collection_with_3_indexes,
         {"listIndexes": collection_with_3_indexes.name, "cursor": {"batchSize": 1}},
     )
-    all_indexes = list(result["cursor"]["firstBatch"])
-    cursor_id = result["cursor"]["id"]
-
-    while cursor_id != 0:
-        get_more_result = execute_command(
-            collection_with_3_indexes,
-            {"getMore": cursor_id, "collection": collection_with_3_indexes.name, "batchSize": 1},
+    if isinstance(result, Exception):
+        outcome = result
+    else:
+        all_indexes = list(result["cursor"]["firstBatch"])
+        cursor_id = result["cursor"]["id"]
+
+        failure = None
+        while cursor_id != 0:
+            get_more_result = execute_command(
+                collection_with_3_indexes,
+                {
+                    "getMore": cursor_id,
+                    "collection": collection_with_3_indexes.name,
+                    "batchSize": 1,
+                },
+            )
+            if isinstance(get_more_result, Exception):
+                failure = get_more_result
+                break
+            all_indexes.extend(get_more_result["cursor"]["nextBatch"])
+            cursor_id = get_more_result["cursor"]["id"]
+        outcome = (
+            failure if failure is not None else {"cursor": {"firstBatch": all_indexes}, "ok": 1.0}
         )
-        all_indexes.extend(get_more_result["cursor"]["nextBatch"])
-        cursor_id = get_more_result["cursor"]["id"]
 
     assertSuccess(
-        {"cursor": {"firstBatch": all_indexes}, "ok": 1.0},
+        outcome,
         [
             {"v": 2, "key": {"_id": 1}, "name": "_id_"},
             {"v": 2, "key": {"a": 1}, "name": "a_1"},
@@ -180,22 +200,36 @@ def test_listIndexes_batchSize_10_with_64_indexes(collection_with_64_indexes):
         collection_with_64_indexes,
         {"listIndexes": collection_with_64_indexes.name, "cursor": {"batchSize": 10}},
     )
-    all_indexes = list(result["cursor"]["firstBatch"])
-    cursor_id = result["cursor"]["id"]
-
-    while cursor_id != 0:
-        get_more_result = execute_command(
-            collection_with_64_indexes,
-            {"getMore": cursor_id, "collection": collection_with_64_indexes.name, "batchSize": 10},
+    if isinstance(result, Exception):
+        outcome = result
+    else:
+        all_indexes = list(result["cursor"]["firstBatch"])
+        cursor_id = result["cursor"]["id"]
+
+        failure = None
+        while cursor_id != 0:
+            get_more_result = execute_command(
+                collection_with_64_indexes,
+                {
+                    "getMore": cursor_id,
+                    "collection": collection_with_64_indexes.name,
+                    "batchSize": 10,
+                },
+            )
+            if isinstance(get_more_result, Exception):
+                failure = get_more_result
+                break
+            all_indexes.extend(get_more_result["cursor"]["nextBatch"])
+            cursor_id = get_more_result["cursor"]["id"]
+        outcome = (
+            failure if failure is not None else {"cursor": {"firstBatch": all_indexes}, "ok": 1.0}
         )
-        all_indexes.extend(get_more_result["cursor"]["nextBatch"])
-        cursor_id = get_more_result["cursor"]["id"]
 
     expected = [{"v": 2, "key": {"_id": 1}, "name": "_id_"}]
     for i in range(63):
         expected.append({"v": 2, "key": {f"f{i}": 1}, "name": f"f{i}_1"})
     assertSuccess(
-        {"cursor": {"firstBatch": all_indexes}, "ok": 1.0},
+        outcome,
         expected,
         msg="Total indexes across all batches should equal 64",
     )

documentdb_tests/compatibility/tests/core/operator/system-stages/changeStream/test_changeStream_event_read_error.py

@@ -44,6 +44,9 @@ def test_changeStream_required_image_event_read_error(collection, test_case):
     opened = execute_command(
         collection, change_stream_command(collection, pipeline=test_case.pipeline)
     )
-    collection.update_one({"_id": 1}, {"$set": {"a": 2}})
-    result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
+    if isinstance(opened, Exception):
+        result = opened
+    else:
+        collection.update_one({"_id": 1}, {"$set": {"a": 2}})
+        result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
     assertResult(result, error_code=test_case.error_code, msg=test_case.msg, raw_res=True)

documentdb_tests/compatibility/tests/core/operator/system-stages/changeStream/test_changeStream_events.py

@@ -425,11 +425,12 @@ def __post_init__(self):
 )
 
 
-def _emit_event(collection, test_case: ChangeStreamEventTestCase) -> dict[str, Any]:
+def _emit_event(collection, test_case: ChangeStreamEventTestCase) -> dict[str, Any] | Exception:
     """Open a stream, perform the test case mutation, and return the matching event.
 
     The mutation happens after the stream is open, so its event is delivered by
-    the first getMore.
+    the first getMore. If opening the stream or reading the batch fails, the
+    engine's exception is returned so the test's single assertion surfaces it.
     """
     mutation = test_case.mutation
     assert mutation is not None  # guaranteed by __post_init__
@@ -443,12 +444,15 @@ def _emit_event(collection, test_case: ChangeStreamEventTestCase) -> dict[str, A
     opened = execute_command(
         collection, change_stream_command(collection, pipeline=test_case.pipeline)
     )
+    if isinstance(opened, Exception):
+        return opened
     mutation.apply(collection)
     if test_case.post_mutation is not None:
         test_case.post_mutation(collection)
-    batch = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))[
-        "cursor"
-    ]["nextBatch"]
+    batch_result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
+    if isinstance(batch_result, Exception):
+        return batch_result
+    batch = batch_result["cursor"]["nextBatch"]
     matching = [e for e in batch if e["operationType"] == mutation.operation_type]
     assert matching, f"no {mutation.operation_type} event in batch: {batch!r}"
     event: dict[str, Any] = matching[0]

documentdb_tests/compatibility/tests/core/operator/system-stages/changeStream/test_changeStream_resume.py

@@ -58,6 +58,36 @@ def _event_token_of_type(collection, operation_type, seed, spec=None):
     return next(e["_id"] for e in batch if e["operationType"] == operation_type)
 
 
+def _first_event_token_or_error(collection, seed, spec=None):
+    """Like ``_first_event_token`` but return the engine error instead of crashing.
+
+    When opening or draining the token-capture stream fails, the exception is
+    returned so a caller can forward it to the test's single assertion rather
+    than indexing into it.
+    """
+    opened = _open_stream(collection, spec)
+    if isinstance(opened, Exception):
+        return opened
+    seed(collection)
+    result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
+    if isinstance(result, Exception):
+        return result
+    return result["cursor"]["nextBatch"][0]["_id"]
+
+
+def _event_token_of_type_or_error(collection, operation_type, seed, spec=None):
+    """Like ``_event_token_of_type`` but return the engine error instead of crashing."""
+    opened = _open_stream(collection, spec)
+    if isinstance(opened, Exception):
+        return opened
+    seed(collection)
+    result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
+    if isinstance(result, Exception):
+        return result
+    batch = result["cursor"]["nextBatch"]
+    return next((e["_id"] for e in batch if e["operationType"] == operation_type), None)
+
+
 def _high_water_mark(collection):
     """Return the postBatchResumeToken of a freshly opened, idle stream."""
     return _open_stream(collection)["cursor"]["postBatchResumeToken"]
@@ -310,14 +340,22 @@ def seed(c):
         c.insert_one({"_id": 1})
         c.drop()
 
-    token = _event_token_of_type(collection, "invalidate", seed)
+    token = _event_token_of_type_or_error(collection, "invalidate", seed)
 
-    opened = execute_command(
-        collection,
-        change_stream_command(collection, pipeline=[{"$changeStream": {"startAfter": token}}]),
-    )
-    collection.insert_one({"_id": 2})
-    result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
+    if isinstance(token, Exception):
+        result = token
+    else:
+        opened = execute_command(
+            collection,
+            change_stream_command(collection, pipeline=[{"$changeStream": {"startAfter": token}}]),
+        )
+        if isinstance(opened, Exception):
+            result = opened
+        else:
+            collection.insert_one({"_id": 2})
+            result = execute_command(
+                collection, get_more_command(collection, opened["cursor"]["id"])
+            )
 
     assertResult(
         result,
@@ -388,14 +426,22 @@ def test_start_after_foreign_collection_token(collection, database_client):
     """Test $changeStream startAfter from a foreign-collection token."""
     other = database_client[f"{collection.name}_other"]
     database_client.create_collection(other.name)
-    token = _first_event_token(other, lambda c: c.insert_one({"_id": 1}))
+    token = _first_event_token_or_error(other, lambda c: c.insert_one({"_id": 1}))
 
-    opened = execute_command(
-        collection,
-        change_stream_command(collection, pipeline=[{"$changeStream": {"startAfter": token}}]),
-    )
-    collection.insert_one({"_id": 99})
-    result = execute_command(collection, get_more_command(collection, opened["cursor"]["id"]))
+    if isinstance(token, Exception):
+        result = token
+    else:
+        opened = execute_command(
+            collection,
+            change_stream_command(collection, pipeline=[{"$changeStream": {"startAfter": token}}]),
+        )
+        if isinstance(opened, Exception):
+            result = opened
+        else:
+            collection.insert_one({"_id": 99})
+            result = execute_command(
+                collection, get_more_command(collection, opened["cursor"]["id"])
+            )
 
     assertResult(
         result,

documentdb_tests/compatibility/tests/core/query_and_write/commands/bulkWrite/test_bulkWrite_response_structure.py

@@ -185,8 +185,14 @@ def test_bulkWrite_getMore_drains_remaining_results(collection):
             "cursor": {"batchSize": 1},
         },
     )
-    cursor_id = result["cursor"]["id"]
-    more = execute_admin_command(collection, {"getMore": cursor_id, "collection": "$cmd.bulkWrite"})
+    cursor_id = None if isinstance(result, Exception) else result["cursor"]["id"]
+    more = (
+        result
+        if isinstance(result, Exception)
+        else execute_admin_command(
+            collection, {"getMore": cursor_id, "collection": "$cmd.bulkWrite"}
+        )
+    )
     assertResult(
         more,
         expected={

documentdb_tests/compatibility/tests/core/query_planning/commands/setQuerySettings/test_setQuerySettings_behavior.py

@@ -133,11 +133,14 @@ def test_setQuerySettings_response(collection, test):
     try:
         result = execute_admin_command(collection, test.build_command(ctx))
         expected = test.build_expected(ctx)
-        # Also verify the dynamic fields are present
-        if test.id == "response_contains_hash":
-            expected["queryShapeHash"] = result.get("queryShapeHash")
-        elif test.id == "response_contains_representative_query":
-            expected["representativeQuery"] = result.get("representativeQuery")
+        # Also verify the dynamic fields are present. When the command failed,
+        # pass the exception through so the single assertion reports the engine
+        # error instead of crashing while reading fields off the exception.
+        if not isinstance(result, Exception):
+            if test.id == "response_contains_hash":
+                expected["queryShapeHash"] = result.get("queryShapeHash")
+            elif test.id == "response_contains_representative_query":
+                expected["representativeQuery"] = result.get("representativeQuery")
         assertSuccessPartial(result, expected, msg=test.msg)
     finally:
         for cmd in test.build_cleanup(ctx):
@@ -231,7 +234,14 @@ def test_setQuerySettings_remove(collection, test):
         for cmd in test.build_setup(ctx):
             r = execute_admin_command(collection, cmd)
             ctx.setup_results.append(r)
-        result = execute_admin_command(collection, test.build_command(ctx))
+        # If a setup command failed, its result is needed to build the command
+        # (e.g. removeQuerySettings by hash); forward the failure to the single
+        # assertion instead of indexing into the exception.
+        setup_failure = next((r for r in ctx.setup_results if isinstance(r, Exception)), None)
+        if setup_failure is not None:
+            result = setup_failure
+        else:
+            result = execute_admin_command(collection, test.build_command(ctx))
         assertSuccessPartial(result, test.build_expected(ctx), msg=test.msg)
     finally:
         for cmd in test.build_cleanup(ctx):

documentdb_tests/compatibility/tests/core/query_planning/commands/setQuerySettings/test_setQuerySettings_settings.py

@@ -775,7 +775,14 @@ def test_setQuerySettings_update(collection, test):
         for cmd in test.build_setup(ctx):
             r = execute_admin_command(collection, cmd)
             ctx.setup_results.append(r)
-        result = execute_admin_command(collection, test.build_command(ctx))
+        # If a setup command failed, its result is needed to build the command
+        # (update via hash); forward the failure to the single assertion instead
+        # of indexing into the exception.
+        setup_failure = next((r for r in ctx.setup_results if isinstance(r, Exception)), None)
+        if setup_failure is not None:
+            result = setup_failure
+        else:
+            result = execute_admin_command(collection, test.build_command(ctx))
         assertSuccessPartial(result, test.build_expected(ctx), msg=test.msg)
     finally:
         for cmd in test.build_cleanup(ctx):

documentdb_tests/compatibility/tests/core/query_planning/commands/setQuerySettings/test_setQuerySettings_verification.py

@@ -326,7 +326,7 @@ def test_setQuerySettings_hash_same(collection, test):
         for cmd in test.build_setup(ctx):
             r = execute_admin_command(collection, cmd)
             ctx.setup_results.append(r)
-            if "queryShapeHash" in r:
+            if isinstance(r, dict) and "queryShapeHash" in r:
                 setup_hash = r["queryShapeHash"]
         result = execute_admin_command(collection, test.build_command(ctx))
         assertSuccessPartial(
@@ -352,12 +352,15 @@ def test_setQuerySettings_hash_different(collection, test):
         for cmd in test.build_setup(ctx):
             r = execute_admin_command(collection, cmd)
             ctx.setup_results.append(r)
-            if "queryShapeHash" in r:
+            if isinstance(r, dict) and "queryShapeHash" in r:
                 setup_hash = r["queryShapeHash"]
         result = execute_admin_command(collection, test.build_command(ctx))
-        hashes_differ = result["queryShapeHash"] != setup_hash
+        if isinstance(result, Exception):
+            actual = result
+        else:
+            actual = {"differ": result["queryShapeHash"] != setup_hash}
         assertSuccessPartial(
-            {"differ": hashes_differ},
+            actual,
             {"differ": True},
             msg=test.msg,
         )
@@ -412,12 +415,17 @@ def test_setQuerySettings_hash_format(collection, test):
     ctx = CommandContext.from_collection(collection)
     try:
         result = execute_admin_command(collection, test.build_command(ctx))
-        h = result.get("queryShapeHash", "")
-        is_valid = bool(re.fullmatch(r"[0-9A-Fa-f]{64}", h))
+        if isinstance(result, Exception):
+            actual = result
+            msg = test.msg
+        else:
+            h = result.get("queryShapeHash", "")
+            actual = {"valid": bool(re.fullmatch(r"[0-9A-Fa-f]{64}", h))}
+            msg = f"{test.msg}, got: {h!r}"
         assertSuccessPartial(
-            {"valid": is_valid},
+            actual,
             {"valid": True},
-            msg=f"{test.msg}, got: {h!r}",
+            msg=msg,
         )
     finally:
         for cmd in test.build_cleanup(ctx):