File size: 3,717 Bytes
69dcc54 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | import pytest
import re
from unittest.mock import patch, MagicMock, AsyncMock
from aioresponses import aioresponses
from app.hunter.strategies.github import GitHubStrategy
from app.hunter.strategies.ai import AIStrategy
from app.hunter.strategies.search import SearchStrategy
@pytest.mark.asyncio
async def test_github_strategy():
strategy = GitHubStrategy()
with aioresponses() as m:
# Mock GitHub Search API
m.get(
re.compile(r"^https://api\.github\.com/search/code.*$"),
payload={
"items": [
{
"html_url": "https://github.com/user/repo/blob/main/proxy.txt",
"name": "proxy.txt",
}
]
},
repeat=True,
)
urls = await strategy.discover()
assert len(urls) > 0
# Check raw conversion
assert "raw.githubusercontent.com" in urls[0]
assert "/blob/" not in urls[0]
@pytest.mark.asyncio
async def test_ai_strategy():
strategy = AIStrategy()
# Create mock response structure
mock_choice = MagicMock()
mock_choice.message.content = (
"Sure! https://pastebin.com/raw/abcd and https://github.com/raw/xyz"
)
mock_response = MagicMock()
mock_response.choices = [mock_choice]
# Mock g4f components
with (
patch("app.hunter.strategies.ai.AsyncClient") as mock_client_cls,
patch("app.hunter.strategies.ai.HAS_G4F", True),
patch("app.hunter.strategies.ai.RetryProvider", MagicMock()),
patch("app.hunter.strategies.ai.PollinationsAI", MagicMock()),
patch("app.hunter.strategies.ai.BlackboxPro", MagicMock()),
):
mock_client = AsyncMock()
mock_client_cls.return_value = mock_client
# Set return value for create
mock_client.chat.completions.create.return_value = mock_response
urls = await strategy.discover()
# Debug info if fails
if len(urls) != 2:
print(f"\nDEBUG: discover returned {urls}")
print(f"DEBUG: AsyncClient patched: {mock_client_cls.called}")
print(f"DEBUG: create called: {mock_client.chat.completions.create.called}")
assert len(urls) == 2
assert "https://pastebin.com/raw/abcd" in urls
assert "https://github.com/raw/xyz" in urls
@pytest.mark.asyncio
async def test_search_strategy():
strategy = SearchStrategy()
# Mock DB storage get_random_proxy
with patch(
"app.hunter.strategies.search.db_storage.get_random_proxy",
new_callable=AsyncMock,
) as mock_get_proxy:
mock_proxy = MagicMock()
mock_proxy.protocol = "http"
mock_proxy.ip = "1.1.1.1"
mock_proxy.port = 8080
mock_get_proxy.return_value = mock_proxy
# Mock aiohttp request to DuckDuckGo
with aioresponses() as m:
# Mock the POST request to html.duckduckgo.com
m.post(
"https://html.duckduckgo.com/html/",
body='<html><body><a class="result__a" href="/l/?kh=-1&uddg=https%3A%2F%2Fpastebin.com%2Fraw%2Ffound">Link</a></body></html>',
repeat=True,
)
# We also need to mock get_db to yield a fake session
with patch("app.hunter.strategies.search.get_db") as mock_get_db:
# Create an async generator mock
async def async_gen():
yield MagicMock()
mock_get_db.return_value = async_gen()
urls = await strategy.discover()
assert len(urls) > 0
assert "https://pastebin.com/raw/found" in urls
|