tests/test_actions/test_google_search.py

import json
from unittest import TestCase, mock

from lagent.actions import GoogleSearch
from lagent.schema import ActionStatusCode


class TestGoogleSearch(TestCase):

    @mock.patch.object(GoogleSearch, '_search')
    def test_search_tool(self, mock_search_func):
        mock_response = (200, json.load('tests/data/search.json'))
        mock_search_func.return_value = mock_response
        search_tool = GoogleSearch(api_key='abc')
        tool_return = search_tool.run("What's the capital of China?")
        self.assertEqual(tool_return.state, ActionStatusCode.SUCCESS)
        self.assertDictEqual(tool_return.result, dict(text="['Beijing']"))

    @mock.patch.object(GoogleSearch, '_search')
    def test_api_error(self, mock_search_func):
        mock_response = (403, {'message': 'bad requests'})
        mock_search_func.return_value = mock_response
        search_tool = GoogleSearch(api_key='abc')
        tool_return = search_tool.run("What's the capital of China?")
        self.assertEqual(tool_return.state, ActionStatusCode.API_ERROR)
        self.assertEqual(tool_return.errmsg, str(403))

    @mock.patch.object(GoogleSearch, '_search')
    def test_http_error(self, mock_search_func):
        mock_response = (-1, 'HTTPSConnectionPool')
        mock_search_func.return_value = mock_response
        search_tool = GoogleSearch(api_key='abc')
        tool_return = search_tool.run("What's the capital of China?")
        self.assertEqual(tool_return.state, ActionStatusCode.HTTP_ERROR)
        self.assertEqual(tool_return.errmsg, 'HTTPSConnectionPool')