From 4f42118d3d4daf9addc1c18f7d6895f7cd74087d Mon Sep 17 00:00:00 2001
From: Mark McDonald <macd@google.com>
Date: Fri, 13 Sep 2024 01:23:40 +0800
Subject: [PATCH] Ensure tests are runnable when standalone (#552)

When we run tests internally that don't invoke the test class, they pass
but don't run the tests.

Our GitHub environment tests them fine but this change ensures they run
in both environments.
---
 tests/test_helpers.py   | 6 +++++-
 tests/test_protos.py    | 5 +++++
 tests/test_responder.py | 7 +++++--
 3 files changed, 15 insertions(+), 3 deletions(-)

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()