Add optional post processing to response recorder

This commit is contained in:
Joakim Soderlund 2018-09-05 19:32:37 +02:00
parent aa6cafa305
commit 3139ff9d7c

View file

@ -26,7 +26,7 @@ import json
from json import JSONDecodeError from json import JSONDecodeError
from os import environ from os import environ
from pathlib import Path from pathlib import Path
from typing import Any, ContextManager, Dict, Iterator, Union, Type from typing import Any, ContextManager, Dict, Iterator, Optional, Union, Type
import importlib_resources as resources import importlib_resources as resources
import pytest import pytest
@ -35,6 +35,8 @@ from requests import Session, Response
from requests.sessions import Request from requests.sessions import Request
from requests_mock import Mocker from requests_mock import Mocker
from fimfarchive.utils import JayWalker
__all__ = ( __all__ = (
'responses', 'responses',
@ -58,6 +60,7 @@ class Recorder(ContextManager['Recorder']):
""" """
self.original = Session.send self.original = Session.send
self.responses: Dict = dict() self.responses: Dict = dict()
self.walker: Optional[JayWalker]
self.path = path self.path = path
def __iter__(self) -> Iterator[Dict[str, Any]]: def __iter__(self) -> Iterator[Dict[str, Any]]:
@ -128,6 +131,10 @@ class Recorder(ContextManager['Recorder']):
Session.send = self.original # type: ignore Session.send = self.original # type: ignore
data = {NAMESPACE: list(self)} data = {NAMESPACE: list(self)}
walker = self.walker
if walker is not None:
walker.walk(data)
with open(self.path, 'wt') as fobj: with open(self.path, 'wt') as fobj:
json.dump(data, fobj, sort_keys=True, indent=4) json.dump(data, fobj, sort_keys=True, indent=4)