Spaces:

Matcry
/

Rabbook

Sleeping

App Files Files Community

Rabbook / tests /test_web_ingest_async.py

Matcry

Deploy snapshot

c76423f 25 days ago

Raw

History Blame Contribute Delete

5.56 kB

	import asyncio
	import unittest
	from unittest.mock import AsyncMock, MagicMock, patch

	from rag.web_ingest import _run_coroutine_blocking, fetch_url_content


	class RunCoroutineBlockingTests(unittest.TestCase):
	"""Tests for the _run_coroutine_blocking helper."""

	def test_returns_result_when_no_loop_running(self):
	# Called from a regular (non-async) context — no event loop is active.
	async def sample():
	return "hello"

	result = _run_coroutine_blocking(sample())
	self.assertEqual(result, "hello")

	def test_returns_result_when_loop_already_running(self):
	# Regression test for the reported bug:
	# asyncio.run() raises "cannot be called from a running event loop" when
	# invoked from inside FastAPI's (or any) already-running loop.
	# _run_coroutine_blocking must handle this without raising.

	async def sample():
	return "from inside loop"

	async def caller():
	# This is the pattern that used to blow up with plain asyncio.run().
	return _run_coroutine_blocking(sample())

	result = asyncio.run(caller())
	self.assertEqual(result, "from inside loop")

	def test_propagates_exceptions(self):
	async def boom():
	raise ValueError("expected error")

	with self.assertRaises(ValueError, msg="expected error"):
	_run_coroutine_blocking(boom())


	class FetchUrlContentAsyncTests(unittest.TestCase):
	"""Tests for fetch_url_content with AsyncWebCrawler mocked."""

	def _make_crawl_result(self, success=True, html="", error_message=""):
	mock_result = MagicMock()
	mock_result.success = success
	mock_result.html = html
	mock_result.error_message = error_message
	return mock_result

	@patch("rag.web_ingest.trafilatura.extract_metadata")
	@patch("rag.web_ingest.trafilatura.extract")
	@patch("rag.web_ingest.AsyncWebCrawler")
	def test_fetch_url_content_returns_page_text(
	self, mock_crawler_cls, mock_extract, mock_extract_metadata
	):
	# Arrange: mock crawler that returns a successful result with HTML
	crawl_result = self._make_crawl_result(
	success=True,
	html="<html><body>Some long article text here</body></html>",
	)
	mock_crawler_instance = AsyncMock()
	mock_crawler_instance.arun = AsyncMock(return_value=crawl_result)
	mock_crawler_cls.return_value.__aenter__ = AsyncMock(
	return_value=mock_crawler_instance
	)
	mock_crawler_cls.return_value.__aexit__ = AsyncMock(return_value=None)

	# trafilatura returns enough text to pass the 200-char threshold
	long_text = "a" * 250
	mock_extract.return_value = long_text
	mock_meta = MagicMock()
	mock_meta.title = "Test Page"
	mock_extract_metadata.return_value = mock_meta

	# Act
	result = fetch_url_content("https://example.com/article")

	# Assert: the page_text from trafilatura is propagated correctly
	self.assertEqual(result["page_text"], long_text)
	self.assertEqual(result["title"], "Test Page")
	self.assertEqual(result["source_url"], "https://example.com/article")
	self.assertEqual(result["domain"], "example.com")
	self.assertTrue(result["file_name"].startswith("url-"))
	self.assertTrue(result["file_name"].endswith(".txt"))

	@patch("rag.web_ingest.AsyncWebCrawler")
	def test_fetch_url_content_raises_on_crawler_failure(self, mock_crawler_cls):
	crawl_result = self._make_crawl_result(
	success=False, error_message="connection refused"
	)
	mock_crawler_instance = AsyncMock()
	mock_crawler_instance.arun = AsyncMock(return_value=crawl_result)
	mock_crawler_cls.return_value.__aenter__ = AsyncMock(
	return_value=mock_crawler_instance
	)
	mock_crawler_cls.return_value.__aexit__ = AsyncMock(return_value=None)

	with self.assertRaises(ValueError, msg="Failed to fetch URL"):
	fetch_url_content("https://example.com/article")

	@patch("rag.web_ingest.trafilatura.extract_metadata")
	@patch("rag.web_ingest.trafilatura.extract")
	@patch("rag.web_ingest.AsyncWebCrawler")
	def test_fetch_url_content_from_running_loop(
	self, mock_crawler_cls, mock_extract, mock_extract_metadata
	):
	# Regression: calling fetch_url_content from inside a running event loop
	# (the FastAPI scenario) must not raise "asyncio.run() cannot be called
	# from a running event loop".
	crawl_result = self._make_crawl_result(success=True, html="<html/>")
	mock_crawler_instance = AsyncMock()
	mock_crawler_instance.arun = AsyncMock(return_value=crawl_result)
	mock_crawler_cls.return_value.__aenter__ = AsyncMock(
	return_value=mock_crawler_instance
	)
	mock_crawler_cls.return_value.__aexit__ = AsyncMock(return_value=None)

	long_text = "b" * 250
	mock_extract.return_value = long_text
	mock_meta = MagicMock()
	mock_meta.title = "Loop Page"
	mock_extract_metadata.return_value = mock_meta

	async def fastapi_route_simulation():
	# Simulates being called from within FastAPI's event loop.
	return fetch_url_content("https://example.com/loop-test")

	result = asyncio.run(fastapi_route_simulation())
	self.assertEqual(result["page_text"], long_text)


	if __name__ == "__main__":
	unittest.main()