tests/test_feedexport_postprocess.py - master - githubmirror/scrapy

tests/test_feedexport_postprocess.py
557 строк19 KB
Adrian
Type tests involving items and exports (#7867)
09 авг 2026, 20:51
09 авг 2026, 20:51a561454
from __future__ import annotations

import bz2
import gzip
import lzma
import marshal
import pickle
import sys
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Any

import pytest

from scrapy.extensions.postprocessing import PostProcessingManager
from scrapy.utils.test import get_crawler
from tests.utils.bases.feedexport import TestFeedExportBase
from tests.utils.decorators import coroutine_test
from tests.utils.feedexport import path_to_url, printf_escape

if TYPE_CHECKING:
    from scrapy import Spider


class TestFeedPostProcessedExports(TestFeedExportBase):
    items = [{"foo": "bar"}]
    expected = b"foo\r\nbar\r\n"

    class MyPlugin1:
        def __init__(self, file, feed_options):
            self.file = file
            self.feed_options = feed_options
            self.char = self.feed_options.get("plugin1_char", b"")

        def write(self, data):
            written_count = self.file.write(data)
            written_count += self.file.write(self.char)
            return written_count

        def close(self):
            self.file.close()

    def _named_tempfile(self, name) -> str:
        return str(Path(self.temp_dir, name))

    async def run_and_export(
        self, spider_cls: type[Spider], settings: dict[str, Any]
    ) -> dict[str, bytes | None]:
        """Run spider with specified settings; return exported data with filename."""

        FEEDS = settings.get("FEEDS") or {}
        settings["FEEDS"] = {
            printf_escape(path_to_url(file_path)): feed_options
            for file_path, feed_options in FEEDS.items()
        }

        content: dict[str, bytes | None] = {}
        try:
            spider_cls.start_urls = [self.mockserver.url("/")]
            crawler = get_crawler(spider_cls, settings)
            await crawler.crawl_async()

            for file_path in FEEDS:
                content[str(file_path)] = (
                    Path(file_path).read_bytes() if Path(file_path).exists() else None
                )

        finally:
            for file_path in FEEDS:
                if not Path(file_path).exists():
                    continue

                Path(file_path).unlink()

        return content

    def get_gzip_compressed(
        self,
        data: bytes,
        compresslevel: int = 9,
        mtime: int = 0,
        filename: str = "",
    ) -> bytes:
        data_stream = BytesIO()
        gzipf = gzip.GzipFile(
            fileobj=data_stream,
            filename=filename,
            mtime=mtime,
            compresslevel=compresslevel,
            mode="wb",
        )
        gzipf.write(data)
        gzipf.close()
        data_stream.seek(0)
        return data_stream.read()

    def test_tell_reports_target_file_position(self):
        """Exporters that wrap the file they get, e.g. through
        :class:`io.TextIOWrapper`, need it to report a position."""
        file = BytesIO()
        manager = PostProcessingManager([self.MyPlugin1], file, {})
        assert manager.tell() == 0
        manager.write(b"foo")
        assert manager.tell() == file.tell() == 3

    @coroutine_test
    async def test_gzip_plugin(self):
        filename = self._named_tempfile("gzip_file")

        settings = {
            "FEEDS": {
                filename: {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                },
            },
        }

        data = await self.exported_data(self.items, settings)
        try:
            gzip.decompress(data[filename])
        except OSError:
            pytest.fail("Received invalid gzip data.")

    @coroutine_test
    async def test_gzip_plugin_compresslevel(self):
        filename_to_compressed = {
            self._named_tempfile("compresslevel_0"): self.get_gzip_compressed(
                self.expected, compresslevel=0
            ),
            self._named_tempfile("compresslevel_9"): self.get_gzip_compressed(
                self.expected, compresslevel=9
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("compresslevel_0"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                    "gzip_compresslevel": 0,
                    "gzip_mtime": 0,
                    "gzip_filename": "",
                },
                self._named_tempfile("compresslevel_9"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                    "gzip_compresslevel": 9,
                    "gzip_mtime": 0,
                    "gzip_filename": "",
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = gzip.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_gzip_plugin_mtime(self):
        filename_to_compressed = {
            self._named_tempfile("mtime_123"): self.get_gzip_compressed(
                self.expected, mtime=123
            ),
            self._named_tempfile("mtime_123456789"): self.get_gzip_compressed(
                self.expected, mtime=123456789
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("mtime_123"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                    "gzip_mtime": 123,
                    "gzip_filename": "",
                },
                self._named_tempfile("mtime_123456789"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                    "gzip_mtime": 123456789,
                    "gzip_filename": "",
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = gzip.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_gzip_plugin_filename(self):
        filename_to_compressed = {
            self._named_tempfile("filename_FILE1"): self.get_gzip_compressed(
                self.expected, filename="FILE1"
            ),
            self._named_tempfile("filename_FILE2"): self.get_gzip_compressed(
                self.expected, filename="FILE2"
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("filename_FILE1"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                    "gzip_mtime": 0,
                    "gzip_filename": "FILE1",
                },
                self._named_tempfile("filename_FILE2"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.GzipPlugin"],
                    "gzip_mtime": 0,
                    "gzip_filename": "FILE2",
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = gzip.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_lzma_plugin(self):
        filename = self._named_tempfile("lzma_file")

        settings = {
            "FEEDS": {
                filename: {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                },
            },
        }

        data = await self.exported_data(self.items, settings)
        try:
            lzma.decompress(data[filename])
        except lzma.LZMAError:
            pytest.fail("Received invalid lzma data.")

    @coroutine_test
    async def test_lzma_plugin_format(self):
        filename_to_compressed = {
            self._named_tempfile("format_FORMAT_XZ"): lzma.compress(
                self.expected, format=lzma.FORMAT_XZ
            ),
            self._named_tempfile("format_FORMAT_ALONE"): lzma.compress(
                self.expected, format=lzma.FORMAT_ALONE
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("format_FORMAT_XZ"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_format": lzma.FORMAT_XZ,
                },
                self._named_tempfile("format_FORMAT_ALONE"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_format": lzma.FORMAT_ALONE,
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = lzma.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_lzma_plugin_check(self):
        filename_to_compressed = {
            self._named_tempfile("check_CHECK_NONE"): lzma.compress(
                self.expected, check=lzma.CHECK_NONE
            ),
            self._named_tempfile("CHECK_SHA256"): lzma.compress(
                self.expected, check=lzma.CHECK_SHA256
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("check_CHECK_NONE"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_check": lzma.CHECK_NONE,
                },
                self._named_tempfile("CHECK_SHA256"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_check": lzma.CHECK_SHA256,
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = lzma.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_lzma_plugin_preset(self):
        filename_to_compressed = {
            self._named_tempfile("preset_PRESET_0"): lzma.compress(
                self.expected, preset=0
            ),
            self._named_tempfile("preset_PRESET_9"): lzma.compress(
                self.expected, preset=9
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("preset_PRESET_0"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_preset": 0,
                },
                self._named_tempfile("preset_PRESET_9"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_preset": 9,
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = lzma.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_lzma_plugin_filters(self):
        if "PyPy" in sys.version:
            # https://foss.heptapod.net/pypy/pypy/-/issues/3527
            pytest.skip("lzma filters doesn't work in PyPy")

        filters = [{"id": lzma.FILTER_LZMA2}]
        compressed = lzma.compress(self.expected, filters=filters)
        filename = self._named_tempfile("filters")

        settings = {
            "FEEDS": {
                filename: {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.LZMAPlugin"],
                    "lzma_filters": filters,
                },
            },
        }

        data = await self.exported_data(self.items, settings)
        assert compressed == data[filename]
        result = lzma.decompress(data[filename])
        assert result == self.expected

    @coroutine_test
    async def test_bz2_plugin(self):
        filename = self._named_tempfile("bz2_file")

        settings = {
            "FEEDS": {
                filename: {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.Bz2Plugin"],
                },
            },
        }

        data = await self.exported_data(self.items, settings)
        try:
            bz2.decompress(data[filename])
        except OSError:
            pytest.fail("Received invalid bz2 data.")

    @coroutine_test
    async def test_bz2_plugin_compresslevel(self):
        filename_to_compressed = {
            self._named_tempfile("compresslevel_1"): bz2.compress(
                self.expected, compresslevel=1
            ),
            self._named_tempfile("compresslevel_9"): bz2.compress(
                self.expected, compresslevel=9
            ),
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("compresslevel_1"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.Bz2Plugin"],
                    "bz2_compresslevel": 1,
                },
                self._named_tempfile("compresslevel_9"): {
                    "format": "csv",
                    "postprocessing": ["scrapy.extensions.postprocessing.Bz2Plugin"],
                    "bz2_compresslevel": 9,
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, compressed in filename_to_compressed.items():
            result = bz2.decompress(data[filename])
            assert compressed == data[filename]
            assert result == self.expected

    @coroutine_test
    async def test_custom_plugin(self):
        filename = self._named_tempfile("csv_file")

        settings = {
            "FEEDS": {
                filename: {
                    "format": "csv",
                    "postprocessing": [self.MyPlugin1],
                },
            },
        }

        data = await self.exported_data(self.items, settings)
        assert data[filename] == self.expected

    @coroutine_test
    async def test_custom_plugin_with_parameter(self):
        expected = b"foo\r\n\nbar\r\n\n"
        filename = self._named_tempfile("newline")

        settings = {
            "FEEDS": {
                filename: {
                    "format": "csv",
                    "postprocessing": [self.MyPlugin1],
                    "plugin1_char": b"\n",
                },
            },
        }

        data = await self.exported_data(self.items, settings)
        assert data[filename] == expected

    @coroutine_test
    async def test_custom_plugin_with_compression(self):
        expected = b"foo\r\n\nbar\r\n\n"

        filename_to_decompressor = {
            self._named_tempfile("bz2"): bz2.decompress,
            self._named_tempfile("lzma"): lzma.decompress,
            self._named_tempfile("gzip"): gzip.decompress,
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("bz2"): {
                    "format": "csv",
                    "postprocessing": [
                        self.MyPlugin1,
                        "scrapy.extensions.postprocessing.Bz2Plugin",
                    ],
                    "plugin1_char": b"\n",
                },
                self._named_tempfile("lzma"): {
                    "format": "csv",
                    "postprocessing": [
                        self.MyPlugin1,
                        "scrapy.extensions.postprocessing.LZMAPlugin",
                    ],
                    "plugin1_char": b"\n",
                },
                self._named_tempfile("gzip"): {
                    "format": "csv",
                    "postprocessing": [
                        self.MyPlugin1,
                        "scrapy.extensions.postprocessing.GzipPlugin",
                    ],
                    "plugin1_char": b"\n",
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, decompressor in filename_to_decompressor.items():
            result = decompressor(data[filename])
            assert result == expected

    @coroutine_test
    async def test_exports_compatibility_with_postproc(self):
        filename_to_expected = {
            self._named_tempfile("csv"): b"foo\r\nbar\r\n",
            self._named_tempfile("json"): b'[\n{"foo": "bar"}\n]',
            self._named_tempfile("jsonlines"): b'{"foo": "bar"}\n',
            self._named_tempfile("xml"): b'<?xml version="1.0" encoding="utf-8"?>\n'
            b"<items>\n<item><foo>bar</foo></item>\n</items>",
        }

        settings = {
            "FEEDS": {
                self._named_tempfile("csv"): {
                    "format": "csv",
                    "postprocessing": [self.MyPlugin1],
                    # empty plugin to activate postprocessing.PostProcessingManager
                },
                self._named_tempfile("json"): {
                    "format": "json",
                    "postprocessing": [self.MyPlugin1],
                },
                self._named_tempfile("jsonlines"): {
                    "format": "jsonlines",
                    "postprocessing": [self.MyPlugin1],
                },
                self._named_tempfile("xml"): {
                    "format": "xml",
                    "postprocessing": [self.MyPlugin1],
                },
                self._named_tempfile("marshal"): {
                    "format": "marshal",
                    "postprocessing": [self.MyPlugin1],
                },
                self._named_tempfile("pickle"): {
                    "format": "pickle",
                    "postprocessing": [self.MyPlugin1],
                },
            },
        }

        data = await self.exported_data(self.items, settings)

        for filename, data_bytes in data.items():
            expected: Any
            result: Any
            if "pickle" in filename:
                expected, result = self.items[0], pickle.loads(data_bytes)
            elif "marshal" in filename:
                expected, result = self.items[0], marshal.loads(data_bytes)
            else:
                expected, result = filename_to_expected[filename], data_bytes
            assert result == expected