Improving coverage

tests/test_transformer.py

@@ -38,13 +38,31 @@ def test_disable_request_header(self):
     def test_disable_response_header(self):
         request = self.factory.get('/')
         headers = {'Content-Type': 'text/html', 'X-Diazo-Off': 'true'}
-        urlopen_mock = get_urlopen_mock()
+        urlopen_mock = get_urlopen_mock(headers=headers)
 
         with patch(URLOPEN, urlopen_mock):
             response = CustomProxyView.as_view()(request, '/')
 
         self.assertEqual(response.content, b'Mock')
 
+    def test_x_diazo_off_false_on_response(self):
+        request = self.factory.get('/')
+        headers = {'Content-Type': 'text/html', 'X-Diazo-Off': 'false'}
+        urlopen_mock = get_urlopen_mock(headers=headers)
+
+        with patch(URLOPEN, urlopen_mock):
+            response = CustomProxyView.as_view()(request, '/')
+
+        self.assertNotIn(response.content, b'Mock')
+
+    def test_x_diazo_off_invalid(self):
+        request = self.factory.get('/')
+        headers = {'Content-Type': 'text/html', 'X-Diazo-Off': 'nopz'}
+        urlopen_mock = get_urlopen_mock(headers=headers)
+
+        with patch(URLOPEN, urlopen_mock), self.assertRaises(ValueError):
+            response = CustomProxyView.as_view()(request, '/')
+
     def test_ajax_request(self):
         request = self.factory.get('/', HTTP_X_REQUESTED_WITH='XMLHttpRequest')