Enable temporary GCS access in tests

PiperOrigin-RevId: 234192557

Author: Ryan Sepassi

oss_scripts/oss_tests.sh

@@ -32,8 +32,26 @@ fi
 TF2_IGNORE=$(for test in $TF2_IGNORE_TESTS; do echo "--ignore=$test "; done)
 
 # Run Tests
-pytest $TF2_IGNORE --ignore="tensorflow_datasets/testing/test_utils.py"
+# Ignores:
+# * Some TF2 tests if running against TF2 (see above)
+# * test_utils.py is not a test file
+# * eager_not_enabled_by_default_test needs to be run separately because the
+#   enable_eager_execution calls set global state and pytest runs all the tests
+#   in the same process.
+pytest \
+  --disable-warnings \
+  $TF2_IGNORE \
+  --ignore="tensorflow_datasets/testing/test_utils.py" \
+  --ignore="tensorflow_datasets/eager_not_enabled_by_default_test.py"
 set_status
+# If not running with TF2, ensure Eager is not enabled by default
+if [[ "$TF_VERSION" != "tf2" ]]
+then
+  pytest \
+    --disable-warnings \
+    tensorflow_datasets/eager_not_enabled_by_default_test.py
+  set_status
+fi
 
 # Test notebooks in isolated environments
 NOTEBOOKS="
@@ -54,8 +72,9 @@ function test_notebook() {
   set_status
 }
 
-# Re-enable as TF 2.0 gets closer to stable release
+# TODO(tfds): Re-enable as TF 2.0 gets closer to stable release
 if [[ "$PY_BIN" = "python2.7" && "$TF_VERSION" = "tf2" ]]
+then
   echo "Skipping notebook tests"
 else
   for notebook in $NOTEBOOKS

tensorflow_datasets/core/dataset_info_test.py

@@ -197,14 +197,17 @@ def test_restore_after_modification(self):
       })
 
   def test_reading_from_gcs_bucket(self):
-    mnist_builder = mnist.MNIST(
-        data_dir=tempfile.mkdtemp(dir=self.get_temp_dir()))
-    info = dataset_info.DatasetInfo(builder=mnist_builder)
-    info = mnist_builder.info
-
-    # A nominal check to see if we read it.
-    self.assertTrue(info.initialized)
-    self.assertEqual(10000, info.splits["test"].num_examples)
+    # The base TestCase prevents GCS access, so we explicitly ask it to restore
+    # access here.
+    with self.gcs_access():
+      mnist_builder = mnist.MNIST(
+          data_dir=tempfile.mkdtemp(dir=self.get_temp_dir()))
+      info = dataset_info.DatasetInfo(builder=mnist_builder)
+      info = mnist_builder.info
+
+      # A nominal check to see if we read it.
+      self.assertTrue(info.initialized)
+      self.assertEqual(10000, info.splits["test"].num_examples)
 
   def test_str_smoke(self):
     info = mnist.MNIST(data_dir="/tmp/some_dummy_dir").info

tensorflow_datasets/testing/test_case.py

@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import contextlib
 import os
 import tempfile
 
@@ -29,6 +30,12 @@
 
 
 
+GCS_ACCESS_FNS = {
+    "original": dataset_info.gcs_dataset_files,
+    "dummy": lambda _: []
+}
+
+
 class TestCase(tf.test.TestCase):
   """Base TestCase to be used for all tests.
 
@@ -41,7 +48,15 @@ def setUpClass(cls):
     super(TestCase, cls).setUpClass()
     cls.test_data = os.path.join(os.path.dirname(__file__), "test_data")
     # Test must not communicate with GCS.
-    dataset_info.gcs_dataset_files = lambda _: []
+    dataset_info.gcs_dataset_files = GCS_ACCESS_FNS["dummy"]
+
+  @contextlib.contextmanager
+  def gcs_access(self):
+    # Restore GCS access
+    dataset_info.gcs_dataset_files = GCS_ACCESS_FNS["original"]
+    yield
+    # Revert access
+    dataset_info.gcs_dataset_files = GCS_ACCESS_FNS["dummy"]
 
   def setUp(self):
     super(TestCase, self).setUp()