mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-08-11 10:07:50 +00:00
36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
|
import os
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from haystack.nodes.search_engine import WebSearch
|
||
|
|
from haystack.schema import Document
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.skipif(
|
||
|
|
not os.environ.get("SERPERDEV_API_KEY", None),
|
||
|
|
reason="Please export an env var called SERPERDEV_API_KEY containing the serper.dev API key to run this test.",
|
||
|
|
)
|
||
|
|
@pytest.mark.integration
|
||
|
|
def test_web_search():
|
||
|
|
ws = WebSearch(api_key=os.environ.get("SERPERDEV_API_KEY", None))
|
||
|
|
result, _ = ws.run(query="Who is the boyfriend of Olivia Wilde?")
|
||
|
|
assert "documents" in result
|
||
|
|
assert len(result["documents"]) > 0
|
||
|
|
assert isinstance(result["documents"][0], Document)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.skipif(
|
||
|
|
not os.environ.get("SERPERDEV_API_KEY", None),
|
||
|
|
reason="Please export an env var called SERPERDEV_API_KEY containing the serper.dev API key to run this test.",
|
||
|
|
)
|
||
|
|
@pytest.mark.integration
|
||
|
|
def test_web_search_with_site_keyword():
|
||
|
|
ws = WebSearch(api_key=os.environ.get("SERPERDEV_API_KEY", None))
|
||
|
|
result, _ = ws.run(query='site:lifewire.com OR site:nasa.gov "electric vehicles"')
|
||
|
|
assert "documents" in result
|
||
|
|
assert len(result["documents"]) > 0
|
||
|
|
assert isinstance(result["documents"][0], Document)
|
||
|
|
assert all(
|
||
|
|
["nasa" in doc.meta["link"] or "lifewire" in doc.meta["link"] for doc in result["documents"]]
|
||
|
|
), "Some documents are not from the specified sites lifewire.com or nasa.gov."
|