Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions hivemind_etl/website/crawlee_client.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import asyncio
from typing import Any
import uuid

from crawlee.playwright_crawler import PlaywrightCrawler, PlaywrightCrawlingContext
from defusedxml import ElementTree as ET
Expand All @@ -21,6 +22,12 @@ def __init__(
# do not persist crawled data to local storage
self.crawler._configuration.persist_storage = False
self.crawler._configuration.write_metadata = False
self.crawler._configuration.purge_on_start = True

# changing the id each time so it wouldn't continue
# fetching the previous links
config = self.crawler._configuration.get_global_configuration()
config.default_request_queue_id = uuid.uuid4().hex

@self.crawler.router.default_handler
async def request_handler(context: PlaywrightCrawlingContext) -> None:
Expand Down
10 changes: 6 additions & 4 deletions tests/unit/test_website_etl.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,17 @@ async def test_extract(self):
"title": "Example",
}
]

# Mock the CrawleeClient class instead of the instance
with patch('hivemind_etl.website.website_etl.CrawleeClient') as MockCrawleeClient:
with patch(
"hivemind_etl.website.website_etl.CrawleeClient"
) as MockCrawleeClient:
mock_client_instance = AsyncMock()
mock_client_instance.crawl.return_value = mocked_data
MockCrawleeClient.return_value = mock_client_instance

extracted_data = await self.website_etl.extract(urls)

self.assertEqual(extracted_data, mocked_data)
MockCrawleeClient.assert_called_once()
mock_client_instance.crawl.assert_awaited_once_with(links=urls)
Expand Down