From 65f4919a363a74d07cc4d3c7de0fb307c88fad15 Mon Sep 17 00:00:00 2001
From: Joakim Soderlund <joakim.soderlund@gmail.com>
Date: Wed, 8 Aug 2018 20:19:37 +0200
Subject: [PATCH] Add fixture for mocking or saving HTTP responses

---
 requirements.txt            |   1 +
 tests/fixtures/__init__.py  |   2 +
 tests/fixtures/responses.py | 182 ++++++++++++++++++++++++++++++++++++
 3 files changed, 185 insertions(+)
 create mode 100644 tests/fixtures/responses.py

diff --git a/requirements.txt b/requirements.txt
index 8850eed..3c59364 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,4 +9,5 @@ jmespath
 mypy
 pytest
 requests
+requests-mock
 tqdm
diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py
index dbb1057..bfa5135 100644
--- a/tests/fixtures/__init__.py
+++ b/tests/fixtures/__init__.py
@@ -23,10 +23,12 @@ Global pytest fixtures.
 
 
 from .common import fetcher, flavor, story
+from .responses import responses
 
 
 __all__ = (
     'fetcher',
     'flavor',
+    'responses',
     'story',
 )
diff --git a/tests/fixtures/responses.py b/tests/fixtures/responses.py
new file mode 100644
index 0000000..0c5276a
--- /dev/null
+++ b/tests/fixtures/responses.py
@@ -0,0 +1,182 @@
+"""
+Requests mocking fixture.
+"""
+
+
+#
+# Fimfarchive, preserves stories from Fimfiction.
+# Copyright (C) 2015  Joakim Soderlund
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+
+
+import json
+from json import JSONDecodeError
+from os import environ
+from pathlib import Path
+from typing import Any, ContextManager, Dict, Iterator, Union, Type
+
+import importlib_resources as resources
+import pytest
+from _pytest.fixtures import FixtureRequest
+from requests import Session, Response
+from requests.sessions import Request
+from requests_mock import Mocker
+
+
+__all__ = (
+    'responses',
+)
+
+
+NAMESPACE = 'responses'
+
+
+class Recorder(ContextManager['Recorder']):
+    """
+    Records responses for mocking.
+    """
+
+    def __init__(self, path: Path) -> None:
+        """
+        Constructor.
+
+        Args:
+            path: File to record to.
+        """
+        self.original = Session.send
+        self.responses: Dict = dict()
+        self.path = path
+
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        """
+        Yields all responses in a JSON-friendly format.
+        """
+        responses = [v for k, v in sorted(self.responses.items())]
+
+        for request, response, exception in responses:
+            if exception is not None:
+                raise exception
+
+            data = {
+                'method': request.method,
+                'status_code': response.status_code,
+                'url': request.url,
+            }
+
+            try:
+                data['json'] = response.json()
+            except JSONDecodeError:
+                data['text'] = response.text
+
+            yield data
+
+    def __call__(
+            self,
+            session: Session,
+            request: Request,
+            **kwargs
+            ) -> Response:
+        """
+        Performs a request and records the response.
+
+        Args:
+            session: The current session.
+            request: The current request.
+            **kwargs: Other arguments.
+
+        Returns:
+            A response instance.
+        """
+        key = (request.url, request.method)
+
+        try:
+            response = self.original(session, request, **kwargs)
+            self.responses[key] = (request, response, None)
+            return response
+        except Exception as exception:
+            self.responses[key] = (request, None, exception)
+            raise exception
+
+    def __enter__(self) -> 'Recorder':
+        """
+        Overrides the session send method.
+        """
+        Session.send = self  # type: ignore
+
+        return self
+
+    def __exit__(self, *args) -> None:
+        """
+        Restores the send method and persists responses.
+        """
+        Session.send = self.original  # type: ignore
+
+        data = {NAMESPACE: list(self)}
+
+        with open(self.path, 'wt') as fobj:
+            json.dump(data, fobj, sort_keys=True, indent=4)
+
+
+class Responder(Mocker, ContextManager['Responder']):
+    """
+    Mocks previously recorded responses.
+    """
+
+    def __init__(self, path: Path) -> None:
+        """
+        Constructor.
+
+        Args:
+            path: File containing the responses.
+        """
+        super().__init__()
+        self.path = path
+
+    def __enter__(self) -> 'Responder':
+        """
+        Enables the responder.
+        """
+        with open(self.path, 'rt') as fobj:
+            data = json.load(fobj)
+
+        mock = super().__enter__()
+        assert mock is self
+
+        for response in data[NAMESPACE]:
+            self.register_uri(**response)
+
+        return self
+
+
+@pytest.fixture(scope='module')
+def responses(request: FixtureRequest) -> Iterator[Union[Recorder, Responder]]:
+    """
+    Mocks or saves HTTP responses.
+    """
+    real = environ.get('REAL_HTTP', '').lower()
+    name = request.fspath.basename[:-3] + '.json'
+    package = request.module.__package__
+
+    context: Type[Union[Recorder, Responder]]
+
+    if real in ('1', 'true', 'yes'):
+        context = Recorder
+    else:
+        context = Responder
+
+    with resources.path(package, name) as path:
+        with context(path) as handler:
+            yield handler