make sure that load users always have the "global" metadata dictionary

and that any request tasks executed by a load user has a metadata dictionary that is the merged version of "global" metadata and any request specific metadata.

where the request specific metadata has precedence over the global one.

grizzly/auth/__init__.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, Literal, Optional, Tuple, Type, TypeVar, Union, cast
 from urllib.parse import urlparse
 
+from grizzly.tasks import RequestTask
 from grizzly.types import GrizzlyResponse
 from grizzly.utils import merge_dicts, safe_del
 
@@ -42,7 +43,7 @@ class AuthType(Enum):
 class GrizzlyHttpAuthClient(Generic[P], metaclass=ABCMeta):
     host: str
     environment: Environment
-    headers: Dict[str, str]
+    metadata: Dict[str, Any]
     cookies: Dict[str, str]
     __context__: ClassVar[Dict[str, Any]] = {
         'verify_certificates': True,
@@ -150,21 +151,34 @@ def refresh_token(client: GrizzlyHttpAuthClient, *args: P.args, **kwargs: P.kwar
                     auth_method = AuthMethod.NONE
 
                 if auth_method is not AuthMethod.NONE and client.session_started is not None:
+                    request: Optional[RequestTask] = None
+                    # look for RequestTask in args, we need to set metadata on it
+                    for arg in args:
+                        if isinstance(arg, RequestTask):
+                            request = arg
+                            break
+
                     session_now = time()
                     session_duration = session_now - client.session_started
 
                     # refresh token if session has been alive for at least refresh_time
-                    if session_duration >= auth_context.get('refresh_time', 3000) or (client.headers.get('Authorization', None) is None and client.cookies == {}):
+                    authorization_token = client.metadata.get('Authorization', None)
+                    if session_duration >= auth_context.get('refresh_time', 3000) or (authorization_token is None and client.cookies == {}):
                         auth_type, secret = self.impl.get_token(client, auth_method)
                         client.session_started = time()
                         if auth_type == AuthType.HEADER:
-                            client.headers.update({'Authorization': f'Bearer {secret}'})
+                            header = {'Authorization': f'Bearer {secret}'}
+                            client.metadata.update(header)
+                            if request is not None:
+                                request.metadata.update(header)
                         else:
                             name, value = secret.split('=', 1)
                             client.cookies.update({name: value})
+                    elif authorization_token is not None and request is not None:  # update current request with active authorization token
+                        request.metadata.update({'Authorization': authorization_token})
                 else:
-                    safe_del(client.headers, 'Authorization')
-                    safe_del(client.headers, 'Cookie')
+                    safe_del(client.metadata, 'Authorization')
+                    safe_del(client.metadata, 'Cookie')
 
             bound = func.__get__(client, client.__class__)
 

grizzly/auth/aad.py

@@ -794,7 +794,7 @@ def get_oauth_token(cls, client: GrizzlyHttpAuthClient, pkcs: Optional[Tuple[str
 
         # build generic header values, but remove stuff that shouldn't be part
         # of authentication flow
-        headers = {**client.headers}
+        headers = {**client.metadata}
         safe_del(headers, 'Authorization')
         safe_del(headers, 'Content-Type')
         safe_del(headers, 'Ocp-Apim-Subscription-Key')

grizzly/events/request_logger.py

@@ -91,6 +91,9 @@ def _get_grizzly_response_user_data(
         request_metadata = request.metadata
         request_payload = request.source
 
+        if request_metadata == {}:
+            request_metadata = None
+
         response_time = kwargs.get('locust_request_meta', {}).get('response_time', None)
 
         stacktrace: Optional[str] = None

grizzly/tasks/clients/http.py

@@ -64,7 +64,7 @@
 @client('http', 'https')
 class HttpClientTask(ClientTask, GrizzlyHttpAuthClient):
     arguments: Dict[str, Any]
-    headers: Dict[str, str]
+    metadata: Dict[str, Any]
     session_started: Optional[float]
     host: str
 
@@ -106,7 +106,7 @@ def __init__(
 
         self.arguments = {}
         self.cookies = {}
-        self.headers = {
+        self.metadata = {
             'x-grizzly-user': f'{self.__class__.__name__}::{id(self)}',
         }
 
@@ -129,9 +129,8 @@ def on_start(self, parent: GrizzlyScenario) -> None:
         self.environment = self.grizzly.state.locust.environment
 
         self.session_started = time()
-        metadata = self._context.get('metadata', None)
-        if metadata is not None:
-            self.headers.update(metadata)
+        metadata = self._context.get('metadata', None) or {}
+        self.metadata.update(metadata)
 
     @refresh_token(AAD)
     def get(self, parent: GrizzlyScenario) -> GrizzlyResponse:
@@ -140,11 +139,11 @@ def get(self, parent: GrizzlyScenario) -> GrizzlyResponse:
 
             meta.update({'request': {
                 'url': url,
-                'metadata': self.headers,
+                'metadata': self.metadata,
                 'payload': None,
             }})
 
-            response = requests.get(url, headers=self.headers, cookies=self.cookies, timeout=30, **self.arguments)
+            response = requests.get(url, headers=self.metadata, cookies=self.cookies, timeout=60, **self.arguments)
 
             payload = response.text
             metadata = dict(response.headers)

grizzly/tasks/request.py

@@ -117,7 +117,7 @@ class RequestTask(GrizzlyMetaRequestTask):
     _template: Optional[Template]
     _source: Optional[str]
     arguments: Optional[Dict[str, str]]
-    metadata: Optional[Dict[str, str]]
+    metadata: Dict[str, str]
     async_request: bool
 
     response: RequestTaskResponse
@@ -129,7 +129,7 @@ def __init__(self, method: RequestMethod, name: str, endpoint: str, source: Opti
         self.name = name
         self.endpoint = endpoint
         self.arguments = None
-        self.metadata = None
+        self.metadata = {}
         self.async_request = False
 
         self._template = None
@@ -186,10 +186,7 @@ def template(self) -> Optional[Template]:
 
     def add_metadata(self, key: str, value: str) -> None:
         """Add new metadata key value, where default value of metadata is None, it must be initialized as a dict."""
-        if self.metadata is None:
-            self.metadata = {key: value}
-        else:
-            self.metadata[key] = value
+        self.metadata.update({key: value})
 
     def __call__(self) -> grizzlytask:
         @grizzlytask

grizzly/users/__init__.py

@@ -63,7 +63,11 @@ def __call__(self, cls: Type[GrizzlyUser]) -> Type[GrizzlyUser]:
         return cls
 
 
-@grizzlycontext(context={'log_all_requests': False, 'variables': {}})
+@grizzlycontext(context={
+    'log_all_requests': False,
+    'variables': {},
+    'metadata': None,
+})
 class GrizzlyUser(User, metaclass=GrizzlyUserMeta):
     __dependencies__: ClassVar[Set[str]] = set()
     __scenario__: GrizzlyContextScenario  # reference to grizzly scenario this user is part of
@@ -83,6 +87,8 @@ class GrizzlyUser(User, metaclass=GrizzlyUserMeta):
     event_hook: GrizzlyEventHook
     grizzly = GrizzlyContext()
 
+    metadata: Dict[str, Any]
+
     def __init__(self, environment: Environment, *args: Any, **kwargs: Any) -> None:
         super().__init__(environment, *args, **kwargs)
 
@@ -93,6 +99,8 @@ def __init__(self, environment: Environment, *args: Any, **kwargs: Any) -> None:
         # these are not copied, and we can share reference
         self._scenario._tasks = self.__scenario__._tasks
 
+        self.metadata = self._context.get('metadata', None) or {}
+
         self.logger = logging.getLogger(f'{self.__class__.__name__}/{id(self)}')
         self.abort = False
         self.event_hook = GrizzlyEventHook()
@@ -146,6 +154,9 @@ def request(self, request: RequestTask) -> GrizzlyResponse:
         start_time = perf_counter()
 
         try:
+            if len(self.metadata or {}) > 0:
+                request.metadata = merge_dicts(self.metadata, request.metadata)
+
             request = self.render(request)
 
             request_impl = self.async_request_impl if isinstance(self, AsyncRequests) and request.async_request else self.request_impl

grizzly/users/restapi.py

@@ -54,6 +54,7 @@
 
 import json
 from abc import ABCMeta
+from copy import copy
 from time import time
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, cast
 
@@ -64,7 +65,7 @@
 
 from grizzly.auth import AAD, GrizzlyHttpAuthClient, refresh_token
 from grizzly.types import GrizzlyResponse, RequestDirection, RequestMethod
-from grizzly.utils import safe_del
+from grizzly.utils import merge_dicts, safe_del
 from grizzly_extras.transformer import TransformerContentType
 
 from . import AsyncRequests, GrizzlyUser, GrizzlyUserMeta, grizzlycontext
@@ -96,30 +97,23 @@ class RestApiUserMeta(GrizzlyUserMeta, ABCMeta):
             'initialize_uri': None,
         },
     },
-    'metadata': None,
 })
 class RestApiUser(GrizzlyUser, AsyncRequests, GrizzlyHttpAuthClient, metaclass=RestApiUserMeta):  # type: ignore[misc]
     session_started: Optional[float]
-    headers: Dict[str, str]
     environment: Environment
 
     timeout: ClassVar[float] = 60.0
 
     def __init__(self, environment: Environment, *args: Any, **kwargs: Any) -> None:
         super().__init__(environment, *args, **kwargs)
 
-        self.headers: Dict[str, str] = {
+        self.metadata = merge_dicts({
             'Content-Type': 'application/json',
             'x-grizzly-user': self.__class__.__name__,
-        }
+        }, self.metadata)
 
         self.session_started = None
 
-        metadata = self._context.get('metadata', None)
-        if metadata is not None:
-            metadata = cast(Dict[str, str], metadata)
-            self.headers.update(metadata)
-
         self.client = FastHttpSession(
             environment=self.environment,
             base_url=self.host,
@@ -182,21 +176,21 @@ def request_impl(self, request: RequestTask) -> GrizzlyResponse:
     @refresh_token(AAD)
     def _request(self, request: RequestTask, client: FastHttpSession) -> GrizzlyResponse:
         """Perform a HTTP request using the provided client. Requests are authenticated if needed."""
+        request_headers = copy(request.metadata or {})
+
         if request.method not in [RequestMethod.GET, RequestMethod.PUT, RequestMethod.POST]:
             message = f'{request.method.name} is not implemented for {self.__class__.__name__}'
             raise NotImplementedError(message)
 
         if request.response.content_type == TransformerContentType.UNDEFINED:
             request.response.content_type = TransformerContentType.JSON
-        elif request.response.content_type == TransformerContentType.XML:
-            self.headers.update({'Content-Type': 'application/xml'})
-        elif request.response.content_type == TransformerContentType.MULTIPART_FORM_DATA:
-            safe_del(self.headers, 'Content-Type')
 
-        if request.metadata is not None:
-            self.headers.update(request.metadata)
+        request_headers.update({'Content-Type': request.response.content_type.value})
+
+        parameters: Dict[str, Any] = {}
 
-        parameters: Dict[str, Any] = {'headers': self.headers}
+        if len(request_headers) > 0:
+            parameters.update({'headers': request_headers})
 
         url = f'{self.host}{request.endpoint}'
 
@@ -235,7 +229,6 @@ def _request(self, request: RequestTask, client: FastHttpSession) -> GrizzlyResp
                 else:
                     message = self._get_error_message(response)
                     message = f'{response.status_code} not in {request.response.status_codes}: {message}'
-                    self.logger.error('%% %r', response._manual_result)
                     response.failure(ResponseError(message))
 
             headers = dict(response.headers.items()) if response.headers not in [None, {}] else None
@@ -252,6 +245,6 @@ 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.headers, 'Authorization')
+            safe_del(self.metadata, 'Authorization')
 
         super().add_context(context)

grizzly_extras/transformer.py

@@ -6,7 +6,6 @@
 
 import re
 from abc import ABCMeta, abstractmethod
-from enum import auto
 from functools import wraps
 from json import JSONEncoder
 from json import dumps as jsondumps
@@ -29,25 +28,19 @@ def __init__(self, message: Optional[str] = None) -> None:
 class TransformerContentType(PermutationEnum):
     __vector__ = (False, True)
 
-    UNDEFINED = 0
-    JSON = auto()
-    XML = auto()
-    PLAIN = auto()
-    MULTIPART_FORM_DATA = auto()
+    UNDEFINED = None
+    JSON = 'application/json'
+    XML = 'application/xml'
+    PLAIN = 'text/plain'
+    MULTIPART_FORM_DATA = 'multipart/form-data'
 
     @classmethod
     def from_string(cls, value: str) -> TransformerContentType:
-        if value.lower().strip() in ['application/json', 'json']:
-            return TransformerContentType.JSON
+        value = value.lower()
 
-        if value.lower().strip() in ['application/xml', 'xml']:
-            return TransformerContentType.XML
-
-        if value.lower().strip() in ['text/plain', 'plain']:
-            return TransformerContentType.PLAIN
-
-        if value.lower().strip() == 'multipart/form-data':
-            return TransformerContentType.MULTIPART_FORM_DATA
+        for enum in cls:
+            if enum.name.lower() == value or enum.value == value:
+                return enum
 
         message = f'"{value}" is an unknown response content type'
         raise ValueError(message)

tests/e2e/test_response_handler_failure.py

@@ -62,7 +62,10 @@ def test_e2e_response_handler_failure(e2e_fixture: End2EndFixture) -> None:
 
     # check request
     assert """metadata:
-<empty>
+{
+  "Content-Type": "application/json",
+  "x-grizzly-user": "RestApiUser_001"
+}
 
 payload:
 {