change context variables during runtime via SetVariableTask

grizzly/auth/__init__.py

@@ -8,7 +8,7 @@
 from functools import wraps
 from importlib import import_module
 from time import time
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, Literal, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, Literal, Optional, Set, Tuple, Type, TypeVar, Union, cast
 from urllib.parse import urlparse
 
 from grizzly.tasks import RequestTask
@@ -64,6 +64,8 @@ class GrizzlyHttpAuthClient(Generic[P], metaclass=ABCMeta):
             },
         },
         'metadata': None,
+        '__cached_auth__': {},
+        '__context_change_history__': set(),
     }
     session_started: Optional[float]
     grizzly: GrizzlyContext
@@ -76,6 +78,14 @@ def add_metadata(self, key: str, value: str) -> None:
 
         cast(dict, self._context['metadata']).update({key: value})
 
+    @property
+    def __context_change_history__(self) -> Set[str]:
+        return cast(Set[str], self._context['__context_change_history__'])
+
+    @property
+    def __cached_auth__(self) -> Dict[str, str]:
+        return cast(Dict[str, str], self._context['__cached_auth__'])
+
 
 AuthenticatableFunc = TypeVar('AuthenticatableFunc', bound=Callable[..., GrizzlyResponse])
 

grizzly/auth/aad.py

@@ -752,7 +752,7 @@ def generate_pkcs() -> Tuple[str, str]:
             request_meta = {
                 'request_type': 'AUTH',
                 'response_time': int((time_perf_counter() - start_time) * 1000),
-                'name': f'{scenario_index} {cls.__name__} OAuth2 user token {version}',
+                'name': f'{scenario_index} {cls.__name__} OAuth2 user token {version}: {username_lowercase}',
                 'context': client._context,
                 'response': None,
                 'exception': exception,
@@ -872,7 +872,7 @@ def get_oauth_token(cls, client: GrizzlyHttpAuthClient, pkcs: Optional[Tuple[str
             request_meta = {
                 'request_type': 'AUTH',
                 'response_time': int((time_perf_counter() - start_time) * 1000),
-                'name': f'{scenario_index} {cls.__name__} OAuth2 user token {version}',
+                'name': f'{scenario_index} {cls.__name__} OAuth2 client token {version}: {auth_client_context["id"]}',
                 'context': client._context,
                 'response': None,
                 'exception': exception,

grizzly/steps/scenario/setup.py

@@ -10,8 +10,9 @@
 from grizzly.auth import GrizzlyHttpAuthClient
 from grizzly.context import GrizzlyContext
 from grizzly.exceptions import RestartScenario
-from grizzly.tasks import GrizzlyTask, RequestTask
+from grizzly.tasks import GrizzlyTask, RequestTask, SetVariableTask
 from grizzly.testdata.utils import create_context_variable, resolve_variable
+from grizzly.types import VariableType
 from grizzly.types.behave import Context, given, register_type, then
 from grizzly.types.locust import StopUser
 from grizzly.utils import is_template, merge_dicts
@@ -31,7 +32,10 @@ def parse_iteration_gramatical_number(text: str) -> str:
 
 @given('set context variable "{variable}" to "{value}"')
 def step_setup_set_context_variable(context: Context, variable: str, value: str) -> None:
-    """Set a variable in the scenario context.
+    """Set a variable in the scenario and user context.
+
+    If this step is before any step that adds a task in the scenario, it will be added to the context which the user is initialized with at start.
+    If it is after any tasks, it will be added as a task which will change the context variable value during runtime.
 
     Variable names can contain (one or more) dot (`.`) or slash (`/`) to indicate that the variable has a nested structure. E.g. `token.url`
     and `token/url` results in `{'token': {'url': '<value'>}}`
@@ -56,9 +60,13 @@ def step_setup_set_context_variable(context: Context, variable: str, value: str)
         value (str): value, data type will be guessed and casted
     """
     grizzly = cast(GrizzlyContext, context.grizzly)
-    context_variable = create_context_variable(grizzly, variable, value)
 
-    grizzly.scenario.context = merge_dicts(grizzly.scenario.context, context_variable)
+    if len(grizzly.scenario.tasks) < 1:
+        context_variable = create_context_variable(grizzly, variable, value)
+
+        grizzly.scenario.context = merge_dicts(grizzly.scenario.context, context_variable)
+    else:
+        grizzly.scenario.tasks.add(SetVariableTask(variable, value, VariableType.CONTEXT))
 
 
 @given('repeat for "{value}" {iteration_number:IterationGramaticalNumber}')

grizzly/steps/setup.py

@@ -10,6 +10,7 @@
 from grizzly.tasks import SetVariableTask
 from grizzly.testdata import GrizzlyVariables
 from grizzly.testdata.utils import resolve_variable
+from grizzly.types import VariableType
 from grizzly.types.behave import Context, given, then
 from grizzly.utils import is_template
 
@@ -118,4 +119,4 @@ def step_setup_variable_value(context: Context, name: str, value: str) -> None:
         except Exception as e:
             raise AssertionError(e) from e
     else:
-        grizzly.scenario.tasks.add(SetVariableTask(name, value))
+        grizzly.scenario.tasks.add(SetVariableTask(name, value, VariableType.VARIABLES))

grizzly/tasks/set_variable.py

@@ -20,6 +20,8 @@
 from typing import TYPE_CHECKING, Any, MutableMapping, Optional, Type, cast
 
 from grizzly.testdata import GrizzlyVariables
+from grizzly.testdata.utils import create_context_variable
+from grizzly.types import VariableType
 from grizzly.utils import is_template
 
 from . import GrizzlyTask, grizzlytask, template
@@ -33,25 +35,30 @@
 class SetVariableTask(GrizzlyTask):
     variable: str
     value: str
+    variable_type: VariableType
 
     _variable_instance: Optional[MutableMapping[str, Any]] = None
     _variable_instance_type: Optional[Type[AtomicVariable]] = None
     _variable_key: str
 
-    def __init__(self, variable: str, value: str) -> None:
+    def __init__(self, variable: str, value: str, variable_type: VariableType) -> None:
         super().__init__()
 
         self.variable = variable
         self.value = value
+        self.variable_type = variable_type
 
-        module_name, variable_type, variable, sub_variable = GrizzlyVariables.get_variable_spec(self.variable)
+        if variable_type == VariableType.VARIABLES:
+            module_name, variable_type_name, variable, sub_variable = GrizzlyVariables.get_variable_spec(self.variable)
 
-        if not (module_name is None or variable_type is None):
-            self._variable_instance_type = GrizzlyVariables.load_variable(module_name, variable_type)
+            if not (module_name is None or variable_type_name is None):
+                self._variable_instance_type = GrizzlyVariables.load_variable(module_name, variable_type_name)
 
-            if not getattr(self._variable_instance_type, '__settable__', False):
-                message = f'{module_name}.{variable_type} is not settable'
-                raise AttributeError(message)
+                if not getattr(self._variable_instance_type, '__settable__', False):
+                    message = f'{module_name}.{variable_type_name} is not settable'
+                    raise AttributeError(message)
+        else:
+            sub_variable = None
 
         if sub_variable is not None:
             self._variable_key = f'{variable}.{sub_variable}'
@@ -61,23 +68,27 @@ def __init__(self, variable: str, value: str) -> None:
     @property
     def variable_template(self) -> str:
         """Create a dummy template for the variable, used so we will not complain that this variable isn't used anywhere."""
-        if not is_template(self.variable):
+        if not is_template(self.variable) and self.variable_type == VariableType.VARIABLES:
             return f'{{{{ {self.variable} }}}}'
 
         return self.variable
 
     def __call__(self) -> grizzlytask:
         @grizzlytask
         def task(parent: GrizzlyScenario) -> Any:
-            if self._variable_instance is None and self._variable_instance_type is not None:
-                self._variable_instance = cast(MutableMapping[str, Any], self._variable_instance_type.get())
-
             value = parent.render(self.value)
 
-            if self._variable_instance is not None:
-                self._variable_instance[self._variable_key] = value
+            if self.variable_type == VariableType.VARIABLES:
+                if self._variable_instance is None and self._variable_instance_type is not None:
+                    self._variable_instance = cast(MutableMapping[str, Any], self._variable_instance_type.get())
+
+
+                if self._variable_instance is not None:
+                    self._variable_instance[self._variable_key] = value
 
-            # always update user context with new value
-            parent.user._context['variables'][self.variable] = value
+                # always update user context with new value
+                parent.user._context['variables'][self.variable] = value
+            else:
+                parent.user.add_context(create_context_variable(self.grizzly, self.variable, value))
 
         return task

grizzly/testdata/utils.py

@@ -15,7 +15,7 @@
 from grizzly.testdata.ast import get_template_variables
 from grizzly.types import GrizzlyVariableType, RequestType, TestdataType
 from grizzly.types.locust import MessageHandler, StopUser
-from grizzly.utils import is_template, merge_dicts
+from grizzly.utils import is_template, merge_dicts, unflatten
 
 from . import GrizzlyVariables
 
@@ -118,12 +118,7 @@ def transform(grizzly: GrizzlyContext, data: Dict[str, Any], scenario: Optional[
 
             paths: List[str] = key.split('.')
             variable = paths.pop(0)
-            path = paths.pop()
-            struct = {path: value}
-            paths.reverse()
-
-            for path in paths:
-                struct = {path: {**struct}}
+            struct = unflatten('.'.join(paths), value)
 
             if variable in testdata:
                 testdata[variable] = merge_dicts(testdata[variable], struct)

grizzly/types/__init__.py

@@ -1,7 +1,7 @@
 """Grizzly types."""
 from __future__ import annotations
 
-from enum import Enum
+from enum import Enum, auto
 from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union, cast
 
 from locust.rpc.protocol import Message
@@ -35,6 +35,11 @@
 ]
 
 
+class VariableType(Enum):
+    VARIABLES = auto()
+    CONTEXT = auto()
+
+
 class MessageDirection(PermutationEnum):
     __vector__ = (True, True)
 

grizzly/users/restapi.py

@@ -36,6 +36,32 @@
 
 See {@pylink grizzly.auth.aad}.
 
+It is possible to change authenticated user during runtime by using {@pylink grizzly.steps.scenario.setup.step_setup_set_context_variable} step expressions inbetween other tasks.
+To change user both `auth.user.username` and `auth.user.password` has to be changed (even though maybe only one of them changes value).
+
+This will then cache the `Authorization` token for the current user, and if changed back to that user there is no need to re-authenticate again, unless `refresh_time` for the first login
+expires.
+
+```gherkin
+Given a user of type "RestApi" load testing "https://api.example.com"
+And repeat for "2" iterations
+And set context variable "user.auth.username" to "bob"
+And set context variable "user.auth.password" to "foobar"
+
+Then get request from endpoint "/api/test"
+
+Given set context variable "user.auth.username" to "alice"
+And set context variable "user.auth.password" to "hello world"
+
+Then get request from endpoint "/api/test"
+
+Given set context variable "user.auth.username" to "bob"
+And set context variable "user.auth.password" to "foobar"
+```
+
+In the above, hypotetical scenario, there will 2 "AAD OAuth2 user token" requests, once for user "bob" and one for user "alice", both done in the first iteration. The second iteration
+the cached authentication tokens will be re-used.
+
 ### Multipart/form-data
 
 RestApi supports posting of multipart/form-data content-type, and in that case additional arguments needs to be passed with the request:
@@ -55,6 +81,7 @@
 import json
 from abc import ABCMeta
 from copy import copy
+from hashlib import sha256
 from html.parser import HTMLParser
 from http.cookiejar import Cookie
 from time import time
@@ -125,6 +152,8 @@ def handle_endtag(self, tag: str) -> None:
             'initialize_uri': None,
         },
     },
+    '__cached_auth__': {},
+    '__context_change_history__': set(),
 })
 class RestApiUser(GrizzlyUser, AsyncRequests, GrizzlyHttpAuthClient, metaclass=RestApiUserMeta):  # type: ignore[misc]
     session_started: Optional[float]
@@ -298,9 +327,69 @@ def _request(self, request: RequestTask, client: FastHttpSession) -> GrizzlyResp
         return (headers, payload)
 
     def add_context(self, context: Dict[str, Any]) -> None:
-        """If context change contains a username we should re-authenticate. This is forced by removing the Authorization header."""
-        # something change in auth context, we need to re-authenticate
-        if context.get('auth', {}).get('user', {}).get('username', None) is not None:
-            safe_del(self.metadata, 'Authorization')
+        """If added context contains changes in `auth`, we should cache current `Authorization` token and force re-auth for a new, if the auth
+        doesn't exist in the cache.
+
+        To force a re-authentication, both auth.user.username and auth.user.password needs be set, even though the actual value is only changed
+        for one of them.
+        """
+        current_username = self._context.get('auth', {}).get('user', {}).get('username', None)
+        current_password = self._context.get('auth', {}).get('user', {}).get('password', None)
+        current_authorization = self.metadata.get('Authorization', None)
+
+        changed_username = context.get('auth', {}).get('user', {}).get('username', None)
+        changed_password = context.get('auth', {}).get('user', {}).get('password', None)
+
+        # check if we're currently have Authorization header connected to a username and password,
+        # and if we're changing either username or password.
+        # if so, we need to cache current username+password token
+        if (
+            current_username is not None
+            and current_password is not None
+            and self.__context_change_history__ == set()
+            and current_authorization is not None
+            and (changed_username is not None or changed_password is not None)
+        ):
+            cache_key_plain = f'{current_username}:{current_password}'
+            cache_key = sha256(cache_key_plain.encode()).hexdigest()
+
+            if cache_key not in self.__cached_auth__:
+                self.__cached_auth__.update({cache_key: current_authorization})
 
         super().add_context(context)
+
+        changed_username_path = 'auth.user.username'
+        changed_password_path = 'auth.user.password'  # noqa: S105
+
+        # update change history if needed
+        context_change_history = self.__context_change_history__
+        if context_change_history != {changed_username_path, changed_password_path}:
+            if changed_username is not None and changed_username_path not in context_change_history:
+                self.__context_change_history__.add(changed_username_path)
+
+            if changed_password is not None and changed_password_path not in context_change_history:
+                self.__context_change_history__.add(changed_password_path)
+
+        # everything needed to force re-auth is not in place
+        if self.__context_change_history__ != {changed_username_path, changed_password_path}:
+            return
+
+        # every context change needed to force re-auth is in place, clear change history
+        self.__context_change_history__.clear()
+
+        username = self._context.get('auth', {}).get('user', {}).get('username', None)
+        password = self._context.get('auth', {}).get('user', {}).get('password', None)
+
+        if username is None or password is None:
+            return
+
+        # check if current username+password has a cached Authorization token
+        cache_key_plain = f'{username}:{password}'
+        cache_key = sha256(cache_key_plain.encode()).hexdigest()
+
+        cached_authorization = self.__cached_auth__.get(cache_key, None)
+
+        if cached_authorization is not None:   # restore from cache
+            self.metadata.update({'Authorization': cached_authorization})
+        else:  # remove from metadata, to force a re-authentication
+            safe_del(self.metadata, 'Authorization')

grizzly/utils.py

@@ -318,6 +318,21 @@ def flatten(node: Dict[str, Any], parents: Optional[List[str]] = None) -> Dict[s
 
     return flat
 
+def unflatten(key: str, value: Any) -> Dict[str, Any]:
+    paths: List[str] = key.split('.')
+
+    # last node should have the value
+    path = paths.pop()
+    struct = {path: value}
+
+    # build the struct from the inside out
+    paths.reverse()
+
+    for path in paths:
+        struct = {path: {**struct}}
+
+    return struct
+
 
 def normalize(value: str) -> str:
     """Normalize a string to make it more non-human friendly."""