scrapy/tests/test_pqueues.py

import tempfile
import unittest

import queuelib

from scrapy.http.request import Request
from scrapy.pqueues import DownloaderAwarePriorityQueue, ScrapyPriorityQueue
from scrapy.spiders import Spider
from scrapy.squeues import FifoMemoryQueue
from scrapy.utils.test import get_crawler
from tests.test_scheduler import MockDownloader, MockEngine


class PriorityQueueTest(unittest.TestCase):
    def setUp(self):
        self.crawler = get_crawler(Spider)
        self.spider = self.crawler._create_spider("foo")

    def test_queue_push_pop_one(self):
        temp_dir = tempfile.mkdtemp()
        queue = ScrapyPriorityQueue.from_crawler(
            self.crawler, FifoMemoryQueue, temp_dir
        )
        self.assertIsNone(queue.pop())
        self.assertEqual(len(queue), 0)
        req1 = Request("https://example.org/1", priority=1)
        queue.push(req1)
        self.assertEqual(len(queue), 1)
        dequeued = queue.pop()
        self.assertEqual(len(queue), 0)
        self.assertEqual(dequeued.url, req1.url)
        self.assertEqual(dequeued.priority, req1.priority)
        self.assertEqual(queue.close(), [])

    def test_no_peek_raises(self):
        if hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
            raise unittest.SkipTest("queuelib.queue.FifoMemoryQueue.peek is defined")
        temp_dir = tempfile.mkdtemp()
        queue = ScrapyPriorityQueue.from_crawler(
            self.crawler, FifoMemoryQueue, temp_dir
        )
        queue.push(Request("https://example.org"))
        with self.assertRaises(
            NotImplementedError,
            msg="The underlying queue class does not implement 'peek'",
        ):
            queue.peek()
        queue.close()

    def test_peek(self):
        if not hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
            raise unittest.SkipTest("queuelib.queue.FifoMemoryQueue.peek is undefined")
        temp_dir = tempfile.mkdtemp()
        queue = ScrapyPriorityQueue.from_crawler(
            self.crawler, FifoMemoryQueue, temp_dir
        )
        self.assertEqual(len(queue), 0)
        self.assertIsNone(queue.peek())
        req1 = Request("https://example.org/1")
        req2 = Request("https://example.org/2")
        req3 = Request("https://example.org/3")
        queue.push(req1)
        queue.push(req2)
        queue.push(req3)
        self.assertEqual(len(queue), 3)
        self.assertEqual(queue.peek().url, req1.url)
        self.assertEqual(queue.pop().url, req1.url)
        self.assertEqual(len(queue), 2)
        self.assertEqual(queue.peek().url, req2.url)
        self.assertEqual(queue.pop().url, req2.url)
        self.assertEqual(len(queue), 1)
        self.assertEqual(queue.peek().url, req3.url)
        self.assertEqual(queue.pop().url, req3.url)
        self.assertEqual(queue.close(), [])

    def test_queue_push_pop_priorities(self):
        temp_dir = tempfile.mkdtemp()
        queue = ScrapyPriorityQueue.from_crawler(
            self.crawler, FifoMemoryQueue, temp_dir, [-1, -2, -3]
        )
        self.assertIsNone(queue.pop())
        self.assertEqual(len(queue), 0)
        req1 = Request("https://example.org/1", priority=1)
        req2 = Request("https://example.org/2", priority=2)
        req3 = Request("https://example.org/3", priority=3)
        queue.push(req1)
        queue.push(req2)
        queue.push(req3)
        self.assertEqual(len(queue), 3)
        dequeued = queue.pop()
        self.assertEqual(len(queue), 2)
        self.assertEqual(dequeued.url, req3.url)
        self.assertEqual(dequeued.priority, req3.priority)
        self.assertEqual(queue.close(), [-1, -2])


class DownloaderAwarePriorityQueueTest(unittest.TestCase):
    def setUp(self):
        crawler = get_crawler(Spider)
        crawler.engine = MockEngine(downloader=MockDownloader())
        self.queue = DownloaderAwarePriorityQueue.from_crawler(
            crawler=crawler,
            downstream_queue_cls=FifoMemoryQueue,
            key="foo/bar",
        )

    def tearDown(self):
        self.queue.close()

    def test_push_pop(self):
        self.assertEqual(len(self.queue), 0)
        self.assertIsNone(self.queue.pop())
        req1 = Request("http://www.example.com/1")
        req2 = Request("http://www.example.com/2")
        req3 = Request("http://www.example.com/3")
        self.queue.push(req1)
        self.queue.push(req2)
        self.queue.push(req3)
        self.assertEqual(len(self.queue), 3)
        self.assertEqual(self.queue.pop().url, req1.url)
        self.assertEqual(len(self.queue), 2)
        self.assertEqual(self.queue.pop().url, req2.url)
        self.assertEqual(len(self.queue), 1)
        self.assertEqual(self.queue.pop().url, req3.url)
        self.assertEqual(len(self.queue), 0)
        self.assertIsNone(self.queue.pop())

    def test_no_peek_raises(self):
        if hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
            raise unittest.SkipTest("queuelib.queue.FifoMemoryQueue.peek is defined")
        self.queue.push(Request("https://example.org"))
        with self.assertRaises(
            NotImplementedError,
            msg="The underlying queue class does not implement 'peek'",
        ):
            self.queue.peek()

    def test_peek(self):
        if not hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
            raise unittest.SkipTest("queuelib.queue.FifoMemoryQueue.peek is undefined")
        self.assertEqual(len(self.queue), 0)
        req1 = Request("https://example.org/1")
        req2 = Request("https://example.org/2")
        req3 = Request("https://example.org/3")
        self.queue.push(req1)
        self.queue.push(req2)
        self.queue.push(req3)
        self.assertEqual(len(self.queue), 3)
        self.assertEqual(self.queue.peek().url, req1.url)
        self.assertEqual(self.queue.pop().url, req1.url)
        self.assertEqual(len(self.queue), 2)
        self.assertEqual(self.queue.peek().url, req2.url)
        self.assertEqual(self.queue.pop().url, req2.url)
        self.assertEqual(len(self.queue), 1)
        self.assertEqual(self.queue.peek().url, req3.url)
        self.assertEqual(self.queue.pop().url, req3.url)
        self.assertIsNone(self.queue.peek())