Merge pull request #1721 from Eshan-Agarwal:patch-33

PiperOrigin-RevId: 302932579

Author: copybara-github

tensorflow_datasets/audio/crema_d_test.py

@@ -40,7 +40,6 @@ class CremaDTest(testing.DatasetBuilderTestCase):
           '1002_AA_ANG_XX.wav', '1003_AA_FEA_XX.wav',
       ]
   }
-  SKIP_CHECKSUMS = True  # TODO(tfds): Update checksums
 
 
 if __name__ == '__main__':

tensorflow_datasets/testing/dataset_builder_testing.py

@@ -148,7 +148,6 @@ def setUpClass(cls):
     tf.enable_v2_behavior()
     super(DatasetBuilderTestCase, cls).setUpClass()
     name = cls.__name__
-    cls._download_urls = set()
     # Check class has the right attributes
     if cls.DATASET_CLASS is None or not callable(cls.DATASET_CLASS):
       raise AssertionError(
@@ -171,6 +170,15 @@ def setUp(self):
     if self.MOCK_OUT_FORBIDDEN_OS_FUNCTIONS:
       self._mock_out_forbidden_os_functions()
 
+    # Track the urls which are downloaded to validate the checksums
+    # The `dl_manager.download` and `dl_manager.download_and_extract` are
+    # patched to record the urls in `_download_urls`.
+    # Calling `dl_manager.download_checksums` stop the url
+    # registration (as checksums are stored remotelly)
+    # `_test_checksums` validates the recorded urls.
+    self._download_urls = set()
+    self._stop_record_download = False
+
   def tearDown(self):
     super(DatasetBuilderTestCase, self).tearDown()
     for patcher in self.patchers:
@@ -222,6 +230,10 @@ def test_info(self):
     self.assertEqual(self.builder.name, info.name)
 
   def _add_url(self, url_or_urls):
+    if self._stop_record_download:
+      # Stop record the checksums if dl_manager.download_checksums has been
+      # called (as checksums may be stored remotelly)
+      return
     if isinstance(url_or_urls, download.resource.Resource):
       self._download_urls.add(url_or_urls.url)
     else:
@@ -244,6 +256,9 @@ def _get_dl_download_result(self, url):
     return utils.map_nested(lambda fname: os.path.join(self.example_dir, fname),
                             self.DL_DOWNLOAD_RESULT)
 
+  def _download_checksums(self, url):
+    self._stop_record_download = True
+
   def _make_builder(self, config=None):
     return self.DATASET_CLASS(  # pylint: disable=not-callable
         data_dir=self.tmp_dir,
@@ -317,7 +332,7 @@ def _download_and_prepare_as_dataset(self, builder):
         "tensorflow_datasets.core.download.DownloadManager",
         download_and_extract=self._get_dl_extract_result,
         download=self._get_dl_download_result,
-        download_checksums=lambda *_: None,
+        download_checksums=self._download_checksums,
         manual_dir=manual_dir,
     ):
       if isinstance(builder, dataset_builder.BeamBasedBuilder):