more ftp status 226 support

3 files changed, 9 insertions, 3 deletions

diff --git a/python/tests/test_grobid.py b/python/tests/test_grobid.py
index 8c5e080..24ea40e 100644
--- a/python/tests/test_grobid.py
+++ b/python/tests/test_grobid.py
@@ -64,8 +64,12 @@ def test_grobid_worker_cdx(grobid_client, wayback_client):
         body=REAL_TEI_XML, content_type='text/xml')
 
     with open('tests/files/example.cdx', 'r') as cdx_file:
-        pusher = CdxLinePusher(worker, cdx_file,
-            filter_http_statuses=[200], filter_mimetypes=['application/pdf'])
+        pusher = CdxLinePusher(
+            worker,
+            cdx_file,
+            filter_http_statuses=[200, 226],
+            filter_mimetypes=['application/pdf'],
+        )
         pusher_counts = pusher.run()
         assert pusher_counts['total']
         assert pusher_counts['pushed'] == 7
diff --git a/python/tests/test_pushers.py b/python/tests/test_pushers.py
index ed9c0bb..52f26c0 100644
--- a/python/tests/test_pushers.py
+++ b/python/tests/test_pushers.py
@@ -19,7 +19,7 @@ def test_cdx_line_pusher():
     # HTTP 200 and application/pdf
     with open('tests/files/example.cdx', 'r') as cdx_file:
         pusher = CdxLinePusher(sink, cdx_file,
-            filter_mimetypes=['application/pdf'], filter_http_statuses=[200])
+            filter_mimetypes=['application/pdf'], filter_http_statuses=[200, 226])
         counts = pusher.run()
     assert counts['total'] == 20
     assert counts['skip-parse'] == 1
diff --git a/python/tests/test_wayback.py b/python/tests/test_wayback.py
index efffbe2..6bc1ca4 100644
--- a/python/tests/test_wayback.py
+++ b/python/tests/test_wayback.py
@@ -119,6 +119,7 @@ def wayback_client(cdx_client, mocker):
     resource = mocker.Mock()
     client.rstore.load_resource = mocker.MagicMock(return_value=resource)
     resource.get_status = mocker.MagicMock(return_value=(200, "Ok"))
+    resource.is_revisit = mocker.MagicMock(return_value=False)
     resource.get_location = mocker.MagicMock(return_value=WARC_TARGET)
     body = mocker.Mock()
     resource.open_raw_content = mocker.MagicMock(return_value=body)
@@ -141,6 +142,7 @@ def wayback_client_pdf(cdx_client, mocker):
     resource = mocker.Mock()
     client.rstore.load_resource = mocker.MagicMock(return_value=resource)
     resource.get_status = mocker.MagicMock(return_value=(200, "Ok"))
+    resource.is_revisit = mocker.MagicMock(return_value=False)
     resource.get_location = mocker.MagicMock(return_value=WARC_TARGET)
     body = mocker.Mock()
     resource.open_raw_content = mocker.MagicMock(return_value=body)