Add support for assertions on raw requests (#1584)

Adds support for assertions on raw requests. Relates #1582 Co-authored-by: Quentin Pradet <[email protected]>
elastic · Sep 27, 2022 · ddad4b4 · ddad4b4
1 parent d97b62b
commit ddad4b4
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 4 deletions.
diff --git a/esrally/driver/runner.py b/esrally/driver/runner.py
@@ -371,13 +371,16 @@ def smaller_than_or_equal(self, expected, actual):
     def equal(self, expected, actual):
         return actual == expected
 
-    def check_assertion(self, op_name, assertion, properties):
+    def check_assertion(self, op_name, assertion, properties, *, props_expand_keys=True):
         path = assertion["property"]
         predicate_name = assertion["condition"]
         expected_value = assertion["value"]
         actual_value = properties
-        for k in path.split("."):
-            actual_value = actual_value[k]
+        if props_expand_keys:
+            for k in path.split("."):
+                actual_value = actual_value[k]
+        else:
+            actual_value = actual_value[path]
         predicate = self.predicates[predicate_name]
         success = predicate(expected_value, actual_value)
         if not success:
@@ -396,6 +399,10 @@ async def __call__(self, *args):
             if isinstance(return_value, dict):
                 for assertion in params["assertions"]:
                     self.check_assertion(op_name, assertion, return_value)
+            elif isinstance(return_value, BytesIO):
+                for assertion in params["assertions"]:
+                    props = parse(return_value, assertion["property"])
+                    self.check_assertion(op_name, assertion, props, props_expand_keys=False)
             else:
                 self.logger.debug("Skipping assertion check in [%s] as [%s] does not return a dict.", op_name, repr(self.delegate))
         return return_value
@@ -1938,7 +1945,7 @@ async def __call__(self, es, params):
             # counter-intuitive, but preserves prior behavior
             headers = None
 
-        await es.perform_request(
+        return await es.perform_request(
             method=params.get("method", "GET"), path=path, headers=headers, body=params.get("body"), params=request_params
         )
 

diff --git a/tests/driver/runner_test.py b/tests/driver/runner_test.py
@@ -232,6 +232,37 @@ async def test_asserts_equal_succeeds(self):
 
         assert final_response == response
 
+    @pytest.mark.asyncio
+    async def test_asserts_text_equal_succeeds(self):
+        es = None
+        response = io.BytesIO(
+            json.dumps(
+                {
+                    "hits": {
+                        "hits": {
+                            "value": 5,
+                            "relation": "eq",
+                        },
+                    },
+                }
+            ).encode()
+        )
+        delegate = mock.AsyncMock(return_value=response)
+        r = runner.AssertingRunner(delegate)
+        async with r:
+            final_response = await r(
+                es,
+                {
+                    "name": "test-task",
+                    "assertions": [
+                        {"property": "hits.hits.value", "condition": "==", "value": 5},
+                        {"property": "hits.hits.relation", "condition": "==", "value": "eq"},
+                    ],
+                },
+            )
+
+        assert final_response == response
+
     @pytest.mark.asyncio
     async def test_asserts_equal_fails(self):
         es = None
@@ -260,6 +291,38 @@ async def test_asserts_equal_fails(self):
                     },
                 )
 
+    @pytest.mark.asyncio
+    async def test_asserts_text_equal_fails(self):
+        es = None
+        response = io.BytesIO(
+            json.dumps(
+                {
+                    "hits": {
+                        "hits": {
+                            "value": 10000,
+                            "relation": "gte",
+                        },
+                    },
+                }
+            ).encode()
+        )
+        delegate = mock.AsyncMock(return_value=response)
+        r = runner.AssertingRunner(delegate)
+        with pytest.raises(
+            exceptions.RallyTaskAssertionError, match=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", "condition": "==", "value": 10000},
+                            {"property": "hits.hits.relation", "condition": "==", "value": "eq"},
+                        ],
+                    },
+                )
+
     @pytest.mark.asyncio
     async def test_skips_asserts_for_non_dicts(self):
         es = None