def test_it_returns_a_task_with_the_injected_headers(self):
locust_request = LocustRequest(method=MagicMock(),
url=MagicMock(),
headers={"x-forwarded-for": ""})
task = Task(name="some name",
request=MagicMock(),
locust_request=locust_request)
expected_headers = {"x-forwarded-for": "1.2.3.4"}
modified_task = Task.inject_headers(task, headers=expected_headers)
assert isinstance(modified_task, Task)
headers = modified_task.locust_request.headers
assert len(headers) == 1
assert headers == expected_headers
def test_it_creates_a_locust_request_when_there_is_none(self):
task = Task(name="some name", request=MagicMock())
modified_task = Task.inject_headers(task, {})
assert modified_task.locust_request
