Merge pull request #13 from SvenskaSpel/add-variable-correlation

Add basic variable correlation feature (only works with json atm)

.correlations

@@ -0,0 +1,3 @@
+# syntax: url for which request to get the data from, regex (with single capture group) to get data from the response, variable name to store the data in, and replace json values in requests
+# this example is a little awkward in that we're using the "token" from a response and putting it in the "job" field in the next request, but I had to make do with what reqres.in offers.
+/api/register,"token":"([^"]*)",job

Makefile

@@ -10,6 +10,7 @@ build:
 up update_tests: 
 	@echo rebuilding har2locust test files
 	har2locust tests/inputs/login.har --plugins har2locust.extra_plugins.plugin_example > tests/outputs/login_plugin.py
+	har2locust tests/inputs/reqres.in.har --plugins har2locust.extra_plugins.plugin_example > tests/outputs/reqres_plugin.py
 	har2locust tests/inputs/login.har --disable-plugins rest.py > tests/outputs/login_disable_rest.py
 	bash -c 'ls tests/inputs/ | xargs -I % basename % .har | xargs -I % bash -c "har2locust tests/inputs/%.har > tests/outputs/%.py"'
 

har2locust/extra_plugins/plugin_example.py

@@ -1,6 +1,7 @@
 # This file has some advanced examples of how to massage your recording
 # Use it as inspiration for the techniques, not as a recommendation for exactly what to do
 import ast
+import pathlib
 import re
 import typing
 
@@ -78,6 +79,47 @@ def visit_Call(self, node: ast.Call) -> ast.Call:
     T().visit(tree)
 
 
+@astprocessor
+def do_correlations(tree: ast.Module, values: dict):
+    correlation_path = pathlib.Path(".correlations")
+    if correlation_path.is_file():
+        with open(correlation_path) as f:
+            correlations = [
+                line.rstrip() for line in f.readlines() if line.strip() and not line.strip().startswith("#")
+            ]
+    else:
+        correlations = []
+
+    class T(ast.NodeTransformer):
+        def visit_Call(self, node: ast.Call) -> ast.Call:
+            for kw in node.keywords:
+                if kw.arg == "json" or kw.arg == "headers":
+                    param_dict: ast.Dict = typing.cast(ast.Dict, kw.value)
+                    for key in param_dict.keys:
+                        for correlation in correlations:
+                            _, _, *corr_vars = correlation.split(",")
+                            if typing.cast(ast.Constant, key).s in corr_vars:
+                                param_dict.values[param_dict.keys.index(key)] = (
+                                    # the variable will always have the name of the first corr var
+                                    ast.Name(corr_vars[0].replace("-", "_"))
+                                )
+
+            self.generic_visit(node)
+            return node
+
+        def visit_With(self, node: ast.With) -> ast.With:
+            for correlation in correlations:
+                url, corr_expr, *corr_vars = correlation.split(",")
+                if url == node.items[0].context_expr.args[1].value:  # type: ignore
+                    node.body[0] = ast.parse(
+                        f"""{corr_vars[0].replace('-', '_')} = re.findall('''{corr_expr}''', resp.text)[0] if resp.text else None"""
+                    )
+            self.generic_visit(node)
+            return node
+
+    T().visit(tree)
+
+
 # More examples
 
 # @entriesprocessor