diff --git a/tests/test_helpers.py b/tests/test_helpers.py index f060caf88..5d9ec1c42 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -12,11 +12,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import pathlib import copy import collections from typing import Union +from absl.testing import absltest from absl.testing import parameterized from google.generativeai import protos @@ -81,3 +81,7 @@ def test_get_model(self, request_options, expected_timeout, expected_retry): self.assertEqual(self.observed_timeout[0], expected_timeout) self.assertEqual(str(self.observed_retry[0]), str(expected_retry)) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/test_protos.py b/tests/test_protos.py index 1b59b0c6e..8a76bd0e3 100644 --- a/tests/test_protos.py +++ b/tests/test_protos.py @@ -15,6 +15,7 @@ import pathlib import re +from absl.testing import absltest from absl.testing import parameterized ROOT = pathlib.Path(__file__).parent.parent @@ -32,3 +33,7 @@ def test_check_glm_imports(self): match, msg=f"Bad `glm.` usage, use `genai.protos` instead,\n in {fpath}", ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/test_responder.py b/tests/test_responder.py index c075fc65a..d2818da8a 100644 --- a/tests/test_responder.py +++ b/tests/test_responder.py @@ -19,8 +19,7 @@ from absl.testing import parameterized from google.generativeai import protos from google.generativeai import responder -import IPython.display -import PIL.Image + HERE = pathlib.Path(__file__).parent TEST_PNG_PATH = HERE / "test_img.png" @@ -250,3 +249,7 @@ def fun(a: annotation): cfd = responder.FunctionDeclaration.from_function(fun) got = cfd.parameters.properties["a"] self.assertEqual(got, expected) + + +if __name__ == "__main__": + absltest.main()