griptape

1
import pytest
2
import os
3
from tests.utils.structure_tester import StructureTester
4

5

6
class TestGoogleDriveClient:
7
    @pytest.fixture(
8
        autouse=True,
9
        params=StructureTester.TOOLKIT_TASK_CAPABLE_PROMPT_DRIVERS,
10
        ids=StructureTester.prompt_driver_id_fn,
11
    )
12
    def structure_tester(self, request):
13
        from griptape.structures import Agent
14
        from griptape.tools import GoogleDriveClient
15

16
        return StructureTester(
17
            Agent(
18
                tools=[
19
                    GoogleDriveClient(
20
                        off_prompt=False,
21
                        service_account_credentials={
22
                            "type": os.environ["GOOGLE_ACCOUNT_TYPE"],
23
                            "project_id": os.environ["GOOGLE_PROJECT_ID"],
24
                            "private_key_id": os.environ["GOOGLE_PRIVATE_KEY_ID"],
25
                            "private_key": os.environ["GOOGLE_PRIVATE_KEY"],
26
                            "client_email": os.environ["GOOGLE_CLIENT_EMAIL"],
27
                            "client_id": os.environ["GOOGLE_CLIENT_ID"],
28
                            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
29
                            "token_uri": "https://oauth2.googleapis.com/token",
30
                            "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
31
                            "client_x509_cert_url": os.environ["GOOGLE_CERT_URL"],
32
                        },
33
                        owner_email=os.environ["GOOGLE_OWNER_EMAIL"],
34
                    )
35
                ],
36
                prompt_driver=request.param,
37
                conversation_memory=None,
38
            )
39
        )
40

41
    def test_list_files(self, structure_tester):
42
        structure_tester.run("List all files on Google Drive.")
43

44
    def test_download_file(self, structure_tester):
45
        structure_tester.run('Download the file called "sample1.txt" from Google Drive.')
46

47
    def test_save_content(self, structure_tester):
48
        structure_tester.run('Save content "Hello, Google Drive!" on Google Drive as "hello.txt".')
49

50
    def test_search_files_by_name(self, structure_tester):
51
        structure_tester.run('Search files with name "hello.txt" on Google Drive.')
52

53
    def test_search_files_by_content(self, structure_tester):
54
        structure_tester.run('Search files with content "Hello, Google Drive!" on Google Drive.')
55