AutoGPT

1
import json
2

3
import pytest
4
from googleapiclient.errors import HttpError
5

6
from autogpt.agents.agent import Agent
7
from autogpt.agents.utils.exceptions import ConfigurationError
8
from autogpt.commands.web_search import google, safe_google_results, web_search
9

10

11
@pytest.mark.parametrize(
12
    "query, expected_output",
13
    [("test", "test"), (["test1", "test2"], '["test1", "test2"]')],
14
)
15
def test_safe_google_results(query, expected_output):
16
    result = safe_google_results(query)
17
    assert isinstance(result, str)
18
    assert result == expected_output
19

20

21
def test_safe_google_results_invalid_input():
22
    with pytest.raises(AttributeError):
23
        safe_google_results(123)
24

25

26
@pytest.mark.parametrize(
27
    "query, num_results, expected_output_parts, return_value",
28
    [
29
        (
30
            "test",
31
            1,
32
            ("Result 1", "https://example.com/result1"),
33
            [{"title": "Result 1", "href": "https://example.com/result1"}],
34
        ),
35
        ("", 1, (), []),
36
        ("no results", 1, (), []),
37
    ],
38
)
39
def test_google_search(
40
    query, num_results, expected_output_parts, return_value, mocker, agent: Agent
41
):
42
    mock_ddg = mocker.Mock()
43
    mock_ddg.return_value = return_value
44

45
    mocker.patch("autogpt.commands.web_search.DDGS.text", mock_ddg)
46
    actual_output = web_search(query, agent=agent, num_results=num_results)
47
    for o in expected_output_parts:
48
        assert o in actual_output
49

50

51
@pytest.fixture
52
def mock_googleapiclient(mocker):
53
    mock_build = mocker.patch("googleapiclient.discovery.build")
54
    mock_service = mocker.Mock()
55
    mock_build.return_value = mock_service
56
    return mock_service.cse().list().execute().get
57

58

59
@pytest.mark.parametrize(
60
    "query, num_results, search_results, expected_output",
61
    [
62
        (
63
            "test",
64
            3,
65
            [
66
                {"link": "http://example.com/result1"},
67
                {"link": "http://example.com/result2"},
68
                {"link": "http://example.com/result3"},
69
            ],
70
            [
71
                "http://example.com/result1",
72
                "http://example.com/result2",
73
                "http://example.com/result3",
74
            ],
75
        ),
76
        ("", 3, [], []),
77
    ],
78
)
79
def test_google_official_search(
80
    query,
81
    num_results,
82
    expected_output,
83
    search_results,
84
    mock_googleapiclient,
85
    agent: Agent,
86
):
87
    mock_googleapiclient.return_value = search_results
88
    actual_output = google(query, agent=agent, num_results=num_results)
89
    assert actual_output == safe_google_results(expected_output)
90

91

92
@pytest.mark.parametrize(
93
    "query, num_results, expected_error_type, http_code, error_msg",
94
    [
95
        (
96
            "invalid query",
97
            3,
98
            HttpError,
99
            400,
100
            "Invalid Value",
101
        ),
102
        (
103
            "invalid API key",
104
            3,
105
            ConfigurationError,
106
            403,
107
            "invalid API key",
108
        ),
109
    ],
110
)
111
def test_google_official_search_errors(
112
    query,
113
    num_results,
114
    expected_error_type,
115
    mock_googleapiclient,
116
    http_code,
117
    error_msg,
118
    agent: Agent,
119
):
120
    class resp:
121
        def __init__(self, _status, _reason):
122
            self.status = _status
123
            self.reason = _reason
124

125
    response_content = {
126
        "error": {"code": http_code, "message": error_msg, "reason": "backendError"}
127
    }
128
    error = HttpError(
129
        resp=resp(http_code, error_msg),
130
        content=str.encode(json.dumps(response_content)),
131
        uri="https://www.googleapis.com/customsearch/v1?q=invalid+query&cx",
132
    )
133

134
    mock_googleapiclient.side_effect = error
135
    with pytest.raises(expected_error_type):
136
        google(query, agent=agent, num_results=num_results)
137