Add save_html_on_codes setting

Author: Mantas Briliauskas

README.rst

@@ -52,9 +52,13 @@ HTML storage downloader middleware supports such options:
 
 * **gzip_output** (bool) - if True, HTML output will be stored in gzip format.
   Default is False.
+* **save_html_on_status** (list) - if not empty, sets list of response codes
+  whitelisted for html saving. If list is empty or not provided, all response
+  codes will be allowed for html saving.
 
 Sample::
 
     HTML_STORAGE = {
-        'gzip_output': True
+        'gzip_output': True,
+        'save_html_on_status': [200, 202]
     }

scrapy_html_storage/__init__.py

@@ -15,6 +15,7 @@ def __init__(self, settings):
         """
         self.settings = settings.get('HTML_STORAGE', {})
         self.gzip_output = self.settings.get('gzip_output', False)
+        self.save_html_on_codes = self.settings.get('save_html_on_codes', [])
 
 
     @classmethod
@@ -43,7 +44,7 @@ def process_response(self, request, response, spider):
         Returns:
             scrapy.http.response.Response: unmodified response object.
         """
-        if should_save_html(request):
+        if self._should_save_html(request, response):
             self._save_html_to(spider.response_html_path(request), response.body)
 
         return response
@@ -64,12 +65,30 @@ def _save_html_to(self, path, html_body):
             fs.write_to_file(path, html_body)
 
 
-def should_save_html(request):
+    def _should_save_html(self, request, response):
+        """
+        Args:
+            request (scrapy.http.request.Request)
+            response (scrapy.http.response.Response)
+
+        Returns:
+            bool: True if this request should be stored to disk, False otherwise.
+        """
+        return 'save_html' in request.meta and \
+            should_save_html_according_response_code(
+                response.status,
+                self.save_html_on_codes
+            )
+
+
+def should_save_html_according_response_code(code, allowed_list):
     """
     Args:
-        request (scrapy.http.request.Request)
+        code (int): response status code
+        allowed_list (list): list of response status codes allowed to save html
 
     Returns:
-        bool: True if this request should be stored to disk, False otherwise.
+        bool: True if allowed_list is empty (save all responses), or response
+              code in allowed list.
     """
-    return 'save_html' in request.meta
+    return not allowed_list or code in allowed_list

tests/test_html_storage_middleware.py

@@ -1,10 +1,10 @@
-from hamcrest import assert_that, is_
+from hamcrest import assert_that, is_, has_properties
 from mock import MagicMock, patch, ANY
 import pytest
 
 from scrapy.settings import Settings
 
-from scrapy_html_storage import HtmlStorageMiddleware, should_save_html
+from scrapy_html_storage import HtmlStorageMiddleware
 
 
 def make_request_mock(save_html=False, query='', results_page=None):
@@ -19,11 +19,51 @@ def make_request_mock(save_html=False, query='', results_page=None):
 
     return request_mock
 
+def make_response_mock(response_status):
+    """ Constructs HTTP Response mock object.
+    """
+    response_mock = MagicMock()
+    response_mock.status = response_status
+
+    return response_mock
+
+
+def make_allowed_response_codes_list():
+    return range(200, 300)
+
+
+def make_downloader(save_html_on_codes=[]):
+    settings = Settings()
+    settings.set('HTML_STORAGE', {
+        'gzip_output': True,
+        'save_html_on_codes': save_html_on_codes
+    })
+    return HtmlStorageMiddleware(settings)
+
+
+@pytest.mark.parametrize('response_status,as_expected', [
+    (200, True),
+    (299, True),
+    (300, False),
+    (404, False),
+])
+def test_should_save_html_returns_true_when_request_metainformation_has_special_key_set_and_appropriate_response_status(response_status, as_expected):
+    request_mock = make_request_mock(save_html=True)
+    response_mock = make_response_mock(response_status=response_status)
+    downloader = make_downloader(make_allowed_response_codes_list())
+
+    save = downloader._should_save_html(request_mock, response_mock)
+
+    assert_that(save, is_(as_expected))
+
 
-def test_should_save_html_returns_true_when_request_metainformation_has_special_key_set():
+@pytest.mark.parametrize('response_status', [200, 299, 300, 404])
+def test_should_save_html_returns_true_when_request_metainformation_has_special_key_set_and_allowed_resonse_codes_list_is_empty(response_status):
     request_mock = make_request_mock(save_html=True)
+    response_mock = make_response_mock(response_status=response_status)
+    downloader = make_downloader()
 
-    save = should_save_html(request_mock)
+    save = downloader._should_save_html(request_mock, response_mock)
 
     assert_that(save, is_(True))
 
@@ -33,8 +73,9 @@ def test_process_response_stores_response_body_to_file_if_request_asks_for_it(
         write_to_file_mock):
     downloader = HtmlStorageMiddleware(Settings())
     request_mock = make_request_mock(save_html=True)
+    response_mock = make_response_mock(response_status=200)
 
-    downloader.process_response(request_mock, MagicMock(), MagicMock())
+    downloader.process_response(request_mock, response_mock, MagicMock())
 
     assert_that(write_to_file_mock.call_count, is_(1))
 
@@ -44,11 +85,12 @@ def test_process_response_saves_response_html_to_file_resolved_by_spider(
         write_to_file_mock):
     downloader = HtmlStorageMiddleware(Settings())
     request_mock = make_request_mock(save_html=True)
+    response_mock = make_response_mock(response_status=200)
 
     spider_mock = MagicMock()
     spider_mock.response_html_path.return_value = '/tmp/response.html'
 
-    downloader.process_response(request_mock, MagicMock(), spider_mock)
+    downloader.process_response(request_mock, response_mock, spider_mock)
 
     write_to_file_mock.assert_called_with('/tmp/response.html', ANY)
 
@@ -59,8 +101,9 @@ def test_process_response_stores_response_body_to_gzip_file_if_this_setting_is_o
     downloader = HtmlStorageMiddleware(Settings())
     downloader.gzip_output = True
     request_mock = make_request_mock(save_html=True)
+    response_mock = make_response_mock(response_status=200)
 
-    downloader.process_response(request_mock, MagicMock(), MagicMock())
+    downloader.process_response(request_mock, response_mock, MagicMock())
 
     assert_that(write_to_gzip_mock.call_count, is_(1))
 
@@ -76,11 +119,18 @@ def test_from_settings_constructs_middleware_with_the_specified_settings():
 
 def test_constructor_extracts_expected_settings():
     settings = Settings()
-    settings.set('HTML_STORAGE', {'gzip_output': True})
+    save_html_on_codes = make_allowed_response_codes_list()
+    settings.set('HTML_STORAGE', {
+        'gzip_output': True,
+        'save_html_on_codes': save_html_on_codes
+    })
 
     downloader = HtmlStorageMiddleware(settings)
 
-    assert_that(downloader.gzip_output, is_(True))
+    assert_that(downloader, has_properties(dict(
+        gzip_output=True,
+        save_html_on_codes=save_html_on_codes
+    )))
 
 
 def test_constructor_sets_empty_settings_when_middleware_settings_are_not_specified():