diff --git a/src/google/adk/tools/bigtable/client.py b/src/google/adk/tools/bigtable/client.py index 204e2756d9..1d1dd417a9 100644 --- a/src/google/adk/tools/bigtable/client.py +++ b/src/google/adk/tools/bigtable/client.py @@ -31,7 +31,7 @@ def _get_client_info() -> google.api_core.client_info.ClientInfo: def get_bigtable_data_client( *, project: str, credentials: Credentials -) -> bigtable.BigtableDataClient: +) -> data.BigtableDataClient: """Get a Bigtable client.""" bigtable_data_client = data.BigtableDataClient( diff --git a/tests/unittests/tools/bigtable/test_client.py b/tests/unittests/tools/bigtable/test_client.py index 533d5ee331..3f48a68874 100644 --- a/tests/unittests/tools/bigtable/test_client.py +++ b/tests/unittests/tools/bigtable/test_client.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import typing from unittest import mock from google.adk.tools.bigtable import client from google.auth.credentials import Credentials +from google.cloud.bigtable import data def test_get_bigtable_data_client(): @@ -48,3 +50,9 @@ def test_get_bigtable_admin_client(): credentials=mock_creds, client_info=mock.ANY, ) + + +def test_get_bigtable_data_client_return_annotation_resolves(): + """The data client return annotation must resolve to a real type.""" + hints = typing.get_type_hints(client.get_bigtable_data_client) + assert hints["return"] is data.BigtableDataClient