From b9dee4002a72f3a08943c8c4b98a3c55d982ef3b Mon Sep 17 00:00:00 2001
From: Daniel Mitterdorfer <daniel@mitterdorfer.name>
Date: Wed, 24 Feb 2021 14:02:48 +0100
Subject: [PATCH] Add optional name in assertion message

With this commit we attempt to extract an operation's `name` property
and if present, we include it in the assertion message. This is mostly
useful for operations inside a composite, as we lack context there and
the error message only indicates an error "somewhere" in the composite
but not exactly in which specific operation. With this commit, it is
possible to get that information.

Closes #1196
---
 esrally/driver/runner.py    | 16 +++++++++++-----
 tests/driver/runner_test.py |  5 ++++-
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/esrally/driver/runner.py b/esrally/driver/runner.py
index e912fe441..1a5e7ff73 100644
--- a/esrally/driver/runner.py
+++ b/esrally/driver/runner.py
@@ -352,7 +352,7 @@ def smaller_than_or_equal(self, expected, actual):
     def equal(self, expected, actual):
         return actual == expected
 
-    def check_assertion(self, assertion, properties):
+    def check_assertion(self, op_name, assertion, properties):
         path = assertion["property"]
         predicate_name = assertion["condition"]
         expected_value = assertion["value"]
@@ -362,18 +362,24 @@ def check_assertion(self, assertion, properties):
         predicate = self.predicates[predicate_name]
         success = predicate(expected_value, actual_value)
         if not success:
-            raise exceptions.RallyTaskAssertionError(
-                f"Expected [{path}] to be {predicate_name} [{expected_value}] but was [{actual_value}].")
+            if op_name:
+                msg = f"Expected [{path}] in [{op_name}] to be {predicate_name} [{expected_value}] but was [{actual_value}]."
+            else:
+                msg = f"Expected [{path}] to be {predicate_name} [{expected_value}] but was [{actual_value}]."
+
+            raise exceptions.RallyTaskAssertionError(msg)
 
     async def __call__(self, *args):
         params = args[1]
         return_value = await self.delegate(*args)
         if AssertingRunner.assertions_enabled and "assertions" in params:
+            op_name = params.get("name")
             if isinstance(return_value, dict):
                 for assertion in params["assertions"]:
-                    self.check_assertion(assertion, return_value)
+                    self.check_assertion(op_name, assertion, return_value)
             else:
-                self.logger.debug("Skipping assertion check as [%s] does not return a dict.", repr(self.delegate))
+                self.logger.debug("Skipping assertion check in [%s] as [%s] does not return a dict.",
+                                  op_name, repr(self.delegate))
         return return_value
 
     def __repr__(self, *args, **kwargs):
diff --git a/tests/driver/runner_test.py b/tests/driver/runner_test.py
index 86dc8b81c..d4d656ac7 100644
--- a/tests/driver/runner_test.py
+++ b/tests/driver/runner_test.py
@@ -162,6 +162,7 @@ async def test_asserts_equal_succeeds(self):
         r = runner.AssertingRunner(delegate)
         async with r:
             final_response = await r(es, {
+                "name": "test-task",
                 "assertions": [
                     {
                         "property": "hits.hits.value",
@@ -193,9 +194,10 @@ async def test_asserts_equal_fails(self):
         delegate.return_value = as_future(response)
         r = runner.AssertingRunner(delegate)
         with self.assertRaisesRegex(exceptions.RallyTaskAssertionError,
-                                    r"Expected \[hits.hits.relation\] to be == \[eq\] but was \[gte\]."):
+                                    r"Expected \[hits.hits.relation\] in \[test-task\] to be == \[eq\] but was \[gte\]."):
             async with r:
                 await r(es, {
+                    "name": "test-task",
                     "assertions": [
                         {
                             "property": "hits.hits.value",
@@ -219,6 +221,7 @@ async def test_skips_asserts_for_non_dicts(self):
         r = runner.AssertingRunner(delegate)
         async with r:
             final_response = await r(es, {
+                "name": "test-task",
                 "assertions": [
                     {
                         "property": "hits.hits.value",