Merge pull request #86 from dusktreader/dusktreader/add_response_comparison_operator

Dusktreader/add response comparison operator

Author: nicoddemus

pytest_flask/plugin.py

@@ -6,6 +6,8 @@
     :copyright: (c) by Vital Kudzelka
     :license: MIT
 """
+import sys
+
 import pytest
 
 from flask import json
@@ -29,6 +31,36 @@ def json(self):
         """
         return json.loads(self.data)
 
+    def __eq__(self, other):
+        if isinstance(other, int):
+            return self.status_code == other
+        # even though the Python 2-specific code works on Python 3, keep the two versions
+        # separate so we can simplify the code once Python 2 support is dropped
+        if sys.version_info[0] == 2:
+            try:
+                super_eq = super(JSONResponse, self).__eq__
+            except AttributeError:
+                return NotImplemented
+            else:
+                return super_eq(other)
+        else:
+            return super(JSONResponse, self).__eq__(other)
+
+    def __ne__(self, other):
+        return not self == other
+
+
+def pytest_assertrepr_compare(op, left, right):
+    if isinstance(left, JSONResponse) and op == '==' and isinstance(right, int):
+        return [
+            'Mismatch in status code for response: {} != {}'.format(
+                left.status_code,
+                right,
+            ),
+            'Response status: {}'.format(left.status),
+        ]
+    return None
+
 
 def _make_test_response_class(response_class):
     """Extends the response class with special attribute to test JSON

tests/test_fixtures.py

@@ -33,6 +33,19 @@ def test_json_response(self, client, accept_json):
         res = client.get(url_for('ping'), headers=accept_json)
         assert res.json == {'ping': 'pong'}
 
+    def test_json_response_compare_to_status_code(self, client, accept_json):
+        assert client.get(url_for('ping'), headers=accept_json) == 200
+        assert client.get('fake-route', headers=accept_json) == 404
+        assert client.get('fake-route', headers=accept_json) != '404'
+        res = client.get(url_for('ping'), headers=accept_json)
+        assert res == res
+
+    def test_mismatching_eq_comparison(self, client, accept_json):
+        with pytest.raises(AssertionError, match=r'Mismatch in status code'):
+            assert client.get('fake-route', headers=accept_json) == 200
+        with pytest.raises(AssertionError, match=r'404 NOT FOUND'):
+            assert client.get('fake-route', headers=accept_json) == '200'
+
     def test_dont_rewrite_existing_implementation(self, app, accept_json):
         class MyResponse(app.response_class):
             @property