path: root/tests/scrapers/test_scraper.py

from hircine.scraper import Scraper, ScrapeWarning


class MockScraper(Scraper):
    is_available = True

    def scrape(self):
        yield lambda: "foo"
        yield "bar"


class NoneScraper(Scraper):
    is_available = True

    def scrape(self):
        yield lambda: "foo"
        yield None


class WarningScraper(Scraper):
    is_available = True

    def warn(self, msg):
        raise ScrapeWarning("Invalid input")

    def scrape(self):
        yield lambda: "foo"
        yield lambda: self.warn("bar")
        yield "baz"


class ParserlessScraper(Scraper):
    is_available = True

    def scrape(self):
        yield "literal"


def test_scraper_collects():
    generator = MockScraper(None).collect()

    assert set(generator) == set(["foo", "bar"])


def test_scraper_collects_with_transformer():
    generator = MockScraper(None).collect([lambda gen, info: map(str.upper, gen)])

    assert set(generator) == set(["FOO", "BAR"])


def test_scraper_collects_warnings():
    scraper = WarningScraper(None)
    generator = scraper.collect()

    assert set(generator) == set(["foo", "baz"])
    assert scraper.get_warnings() == ["Invalid input"]


def test_scraper_collects_literal():
    scraper = ParserlessScraper(None)
    generator = scraper.collect()

    assert set(generator) == set(["literal"])


def test_scraper_collect_ignores_none():
    scraper = NoneScraper(None)
    generator = scraper.collect()

    assert set(generator) == set(["foo"])