diff --git a/llama-index-integrations/readers/llama-index-readers-google/CHANGELOG.md b/llama-index-integrations/readers/llama-index-readers-google/CHANGELOG.md
index 34e9ee4085e9cfc32fcfe7630c9886f4cfdfc516..0dbc340e52f326f4454b36f2867427b0dedb5c38 100644
--- a/llama-index-integrations/readers/llama-index-readers-google/CHANGELOG.md
+++ b/llama-index-integrations/readers/llama-index-readers-google/CHANGELOG.md
@@ -1,5 +1,9 @@
# CHANGELOG
+## [0.2.2] - 2024-03-26
+
+- Add class name method for serialization
+
## [0.2.1] - 2024-03-26
- Allow passing credentials directly as a string
diff --git a/llama-index-integrations/readers/llama-index-readers-google/llama_index/readers/google/drive/base.py b/llama-index-integrations/readers/llama-index-readers-google/llama_index/readers/google/drive/base.py
index 96c9b6750882782cebe535ec2f1355a4e908724b..d07a2de0b2341527aa12015415541b36d8b58503 100644
--- a/llama-index-integrations/readers/llama-index-readers-google/llama_index/readers/google/drive/base.py
+++ b/llama-index-integrations/readers/llama-index-readers-google/llama_index/readers/google/drive/base.py
@@ -44,8 +44,6 @@ class GoogleDriveReader(BasePydanticReader):
user info. Defaults to None.
service_account_key (Optional[dict]): Dictionary containing service
account key. Defaults to None.
-
-
"""
client_config: Optional[dict] = None
@@ -114,6 +112,10 @@ class GoogleDriveReader(BasePydanticReader):
**kwargs,
)
+ @classmethod
+ def class_name(cls) -> str:
+ return "GoogleDriveReader"
+
def _get_credentials(self) -> Tuple[Credentials]:
"""Authenticate with Google and save credentials.
Download the service_account_key.json file with these instructions: https://cloud.google.com/iam/docs/keys-create-delete.
diff --git a/llama-index-integrations/readers/llama-index-readers-google/pyproject.toml b/llama-index-integrations/readers/llama-index-readers-google/pyproject.toml
index 8c858b15e921e16bd8902404294369105fd862d1..046213bbd4f6116408bf52bce7c3646380109dbf 100644
--- a/llama-index-integrations/readers/llama-index-readers-google/pyproject.toml
+++ b/llama-index-integrations/readers/llama-index-readers-google/pyproject.toml
@@ -45,7 +45,7 @@ maintainers = [
]
name = "llama-index-readers-google"
readme = "README.md"
-version = "0.2.1"
+version = "0.2.2"
[tool.poetry.dependencies]
python = ">=3.10,<4.0"
diff --git a/llama-index-integrations/readers/llama-index-readers-google/tests/test_drive_id.py b/llama-index-integrations/readers/llama-index-readers-google/tests/test_drive_id.py
deleted file mode 100644
index 6ffc4e921fa7bdbfd8ba08618cd64ba6c0fff9ce..0000000000000000000000000000000000000000
--- a/llama-index-integrations/readers/llama-index-readers-google/tests/test_drive_id.py
+++ /dev/null
@@ -1,29 +0,0 @@
-import unittest
-from unittest.mock import MagicMock
-from llama_index.readers.google import GoogleDriveReader
-
-
-class TestGoogleDriveReader(unittest.TestCase):
- def test_load_data_with_drive_id(self):
- # Mock the necessary objects and methods
- mock_credentials = MagicMock()
- mock_drive = MagicMock()
- reader = GoogleDriveReader()
- reader._get_credentials = MagicMock(return_value=(mock_credentials, mock_drive))
- reader._load_from_folder = MagicMock(return_value=["document1", "document2"])
-
- # Test with a specific drive_id
- drive_id = "example_drive_id"
- folder_id = "example_folder_id"
- result = reader.load_data(drive_id=drive_id, folder_id=folder_id)
-
- # Assert that the correct methods are called and the correct result is returned
- reader._get_credentials.assert_called_once()
- reader._load_from_folder.assert_called_once_with(
- drive_id, folder_id, None, None
- )
- self.assertEqual(result, ["document1", "document2"])
-
-
-if __name__ == "__main__":
- unittest.main()
diff --git a/llama-index-integrations/readers/llama-index-readers-google/tests/test_readers_google_drive.py b/llama-index-integrations/readers/llama-index-readers-google/tests/test_readers_google_drive.py
index 871f4ca7ad87eff2884f7006d13013061504fde7..23ccdaf2f3c78a455260158eec3a4fe23ddff1c3 100644
--- a/llama-index-integrations/readers/llama-index-readers-google/tests/test_readers_google_drive.py
+++ b/llama-index-integrations/readers/llama-index-readers-google/tests/test_readers_google_drive.py
@@ -1,6 +1,10 @@
import json
from tempfile import TemporaryDirectory
+import unittest
+from unittest.mock import MagicMock
+from llama_index.readers.google import GoogleDriveReader
+
from llama_index.core.readers.base import BaseReader
import pytest
from llama_index.readers.google import GoogleDriveReader
@@ -10,51 +14,79 @@ test_authorized_user_info = {"authorized_user_info": {"key": "value"}}
test_service_account_key = {"service_account_key": {"key": "value"}}
-def test_class():
- names_of_base_classes = [b.__name__ for b in GoogleDriveReader.__mro__]
- assert BaseReader.__name__ in names_of_base_classes
+class TestGoogleDriveReader(unittest.TestCase):
+ def test_class(self):
+ names_of_base_classes = [b.__name__ for b in GoogleDriveReader.__mro__]
+ assert BaseReader.__name__ in names_of_base_classes
+ def test_serialize(self):
+ reader = GoogleDriveReader(
+ client_config=test_client_config,
+ authorized_user_info=test_authorized_user_info,
+ )
-def test_serialize():
- reader = GoogleDriveReader(
- client_config=test_client_config,
- authorized_user_info=test_authorized_user_info,
- )
+ schema = reader.schema()
+ assert schema is not None
+ assert len(schema) > 0
+ assert "client_config" in schema["properties"]
- schema = reader.schema()
- assert schema is not None
- assert len(schema) > 0
- assert "client_config" in schema["properties"]
+ json_reader = reader.json(exclude_unset=True)
- json_reader = reader.json(exclude_unset=True)
+ new_reader = GoogleDriveReader.parse_raw(json_reader)
+ assert new_reader.client_config == reader.client_config
+ assert new_reader.authorized_user_info == reader.authorized_user_info
- new_reader = GoogleDriveReader.parse_raw(json_reader)
- assert new_reader.client_config == reader.client_config
- assert new_reader.authorized_user_info == reader.authorized_user_info
+ def test_serialize_from_file(self):
+ with TemporaryDirectory() as tmp_dir:
+ file_name = f"{tmp_dir}/test_service_account_key.json"
+ with open(file_name, "w") as f:
+ f.write(json.dumps(test_service_account_key))
+ reader = GoogleDriveReader(
+ service_account_key_path=file_name,
+ )
-def test_serialize_from_file():
- with TemporaryDirectory() as tmp_dir:
- file_name = f"{tmp_dir}/test_service_account_key.json"
- with open(file_name, "w") as f:
- f.write(json.dumps(test_service_account_key))
+ schema = reader.schema()
+ assert schema is not None
+ assert len(schema) > 0
+ assert "service_account_key" in schema["properties"]
- reader = GoogleDriveReader(
- service_account_key_path=file_name,
- )
+ json_reader = reader.json(exclude_unset=True)
- schema = reader.schema()
- assert schema is not None
- assert len(schema) > 0
- assert "service_account_key" in schema["properties"]
+ new_reader = GoogleDriveReader.parse_raw(json_reader)
+ assert new_reader.service_account_key == reader.service_account_key
+ assert new_reader.service_account_key == test_service_account_key
- json_reader = reader.json(exclude_unset=True)
+ def test_error_on_missing_args(self):
+ with pytest.raises(ValueError):
+ GoogleDriveReader()
- new_reader = GoogleDriveReader.parse_raw(json_reader)
- assert new_reader.service_account_key == reader.service_account_key
- assert new_reader.service_account_key == test_service_account_key
+ def test_load_data_with_drive_id(self):
+ # Mock the necessary objects and methods
+ mock_credentials = MagicMock()
+ mock_drive = MagicMock()
+ GoogleDriveReader._get_credentials = MagicMock(
+ return_value=(mock_credentials, mock_drive)
+ )
+ GoogleDriveReader._load_from_folder = MagicMock(
+ return_value=["document1", "document2"]
+ )
+ reader = GoogleDriveReader(
+ client_config={
+ "client_id": "example_client_id",
+ "client_secret": "example_client_secret",
+ },
+ )
+
+ # Test with a specific drive_id
+ drive_id = "example_drive_id"
+ folder_id = "example_folder_id"
+ result = reader.load_data(drive_id=drive_id, folder_id=folder_id)
+ # Assert that the correct methods are called and the correct result is returned
+ reader._get_credentials.assert_called_once()
+ reader._load_from_folder.assert_called_once_with(
+ drive_id, folder_id, None, None
+ )
-def test_error_on_missing_args():
- with pytest.raises(ValueError):
- GoogleDriveReader()
+ assert result == ["document1", "document2"]