haystack/test/test_connector.py

import json
from pathlib import Path
from re import search

import pytest
from haystack.nodes.connector import Crawler


def test_crawler_url_none_exception(tmp_path):
    tmp_dir = tmp_path / "crawled_files"
    with pytest.raises(ValueError):
        Crawler(tmp_dir).crawl()


def test_crawler_depth(tmp_path):
    tmp_dir = tmp_path / "crawled_files"
    _url = ["https://haystack.deepset.ai/overview/get-started"]
    crawler = Crawler(output_dir=tmp_dir)
    doc_path = crawler.crawl(urls=_url, crawler_depth=0)
    assert len(doc_path) == 1

    _urls = [
        "https://haystack.deepset.ai/overview/v0.8.0/get-started",
        "https://haystack.deepset.ai/overview/v0.7.0/get-started",
        "https://haystack.deepset.ai/overview/v0.6.0/get-started",
    ]
    doc_path = crawler.crawl(urls=_urls, crawler_depth=0)
    assert len(doc_path) == 3

    doc_path = crawler.crawl(urls=_url, crawler_depth=1)
    assert len(doc_path) > 1

    for json_file in doc_path:
        assert isinstance(json_file, Path)
        with open(json_file.absolute(), "r") as read_file:
            data = json.load(read_file)
            assert 'content' in data
            assert 'meta' in data
            assert isinstance(data['content'], str)
            assert len(data['content'].split()) > 2


def test_crawler_filter_urls(tmp_path):
    tmp_dir = tmp_path / "crawled_files"
    _url = ["https://haystack.deepset.ai/overview/v0.8.0/"]

    crawler = Crawler(output_dir=tmp_dir)
    doc_path = crawler.crawl(urls=_url, filter_urls=["haystack\.deepset\.ai\/overview\/v0\.9\.0\/"])
    assert len(doc_path) == 0

    doc_path = crawler.crawl(urls=_url, filter_urls=["haystack\.deepset\.ai\/overview\/v0\.8\.0\/"])
    assert len(doc_path) > 0

    doc_path = crawler.crawl(urls=_url, filter_urls=["google\.com"])
    assert len(doc_path) == 0


def test_crawler_content(tmp_path):
    tmp_dir = tmp_path / "crawled_files"

    partial_content_match: list = [
        {"url": "https://haystack.deepset.ai/overview/v0.7.0/intro",
         "partial_content": ["What is Haystack",
                             "Utilize all transformer based models",
                             "a Retriever-Reader pipeline in order",
                             "Passing on only a small candidate set",
                             "fast indexing and querying",
                             "Fine-tune models to your own domain",
                             "smoothly switch when new ones get published"]},
        {"url": "https://haystack.deepset.ai/overview/v0.7.0/use-cases",
         "partial_content": ["Semantic Search System",
                             "Store your documents in the database of ",
                             "results are chosen based on compatibility in",
                             "Apply a set of standard questions to each document",
                             "Return a NO_ANSWER if a given document",
                             "like what is the revenue forecast for 2020?",
                             "overview of academic papers and internal business"]}]

    crawler = Crawler(output_dir=tmp_dir)
    for _dict in partial_content_match:
        url: str = _dict['url']
        partial_content: list = _dict['partial_content']

        doc_path = crawler.crawl(urls=[url], crawler_depth=0)
        assert len(doc_path) == 1

        for json_file in doc_path:
            assert isinstance(json_file, Path)
            with open(json_file.absolute(), "r") as read_file:
                content = json.load(read_file)
                assert isinstance(content['content'], str)
                for partial_line in partial_content:
                    assert search(partial_line, content['content'])
                    assert partial_line in content['content']


def test_crawler_return_document(tmp_path):
    tmp_dir = tmp_path / "crawled_files"
    _url = ["https://haystack.deepset.ai/docs/v0.5.0/intromd"]

    crawler = Crawler(output_dir=tmp_dir)
    docs_path = crawler.crawl(urls=_url, crawler_depth=1)
    results, _ = crawler.run(urls=_url, crawler_depth=1, return_documents=True)
    documents = results['documents']

    for json_file, document in zip(docs_path, documents):
        assert isinstance(json_file, Path)
        assert isinstance(document, dict)

        with open(json_file.absolute(), "r") as read_file:
            file_content = json.load(read_file)
            assert file_content['meta'] == document['meta']
            assert file_content['content'] == document['content']