diff --git a/tests/implicitron/test_data_source.py b/tests/implicitron/test_data_source.py index 950e0c10..f289fef2 100644 --- a/tests/implicitron/test_data_source.py +++ b/tests/implicitron/test_data_source.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import contextlib import os import unittest import unittest.mock @@ -18,6 +19,7 @@ from pytorch3d.implicitron.dataset.data_source import ImplicitronDataSource from pytorch3d.implicitron.dataset.json_index_dataset import JsonIndexDataset from pytorch3d.implicitron.tools.config import get_default_args from tests.common_testing import get_tests_dir +from tests.implicitron.common_resources import get_skateboard_data DATA_DIR = get_tests_dir() / "implicitron/data" DEBUG: bool = False @@ -28,6 +30,12 @@ class TestDataSource(unittest.TestCase): self.maxDiff = None torch.manual_seed(42) + stack = contextlib.ExitStack() + self.dataset_root, self.path_manager = stack.enter_context( + get_skateboard_data() + ) + self.addCleanup(stack.close) + def _test_omegaconf_generic_failure(self): # OmegaConf possible bug - this is why we need _GenericWorkaround from dataclasses import dataclass @@ -56,12 +64,14 @@ class TestDataSource(unittest.TestCase): get_default_args(JsonIndexDataset) def test_one(self): - with unittest.mock.patch.dict(os.environ, {"CO3D_DATASET_ROOT": ""}): - cfg = get_default_args(ImplicitronDataSource) - yaml = OmegaConf.to_yaml(cfg, sort_keys=False) - if DEBUG: - (DATA_DIR / "data_source.yaml").write_text(yaml) - self.assertEqual(yaml, (DATA_DIR / "data_source.yaml").read_text()) + cfg = get_default_args(ImplicitronDataSource) + # making the test invariant to env variables + cfg.dataset_map_provider_JsonIndexDatasetMapProvider_args.dataset_root = "" + cfg.dataset_map_provider_JsonIndexDatasetMapProviderV2_args.dataset_root = "" + yaml = OmegaConf.to_yaml(cfg, sort_keys=False) + if DEBUG: + (DATA_DIR / "data_source.yaml").write_text(yaml) + self.assertEqual(yaml, (DATA_DIR / "data_source.yaml").read_text()) def test_default(self): if os.environ.get("INSIDE_RE_WORKER") is not None: @@ -73,7 +83,7 @@ class TestDataSource(unittest.TestCase): dataset_args.test_restrict_sequence_id = 0 dataset_args.n_frames_per_sequence = -1 - dataset_args.dataset_root = "manifold://co3d/tree/extracted" + dataset_args.dataset_root = self.dataset_root data_source = ImplicitronDataSource(**args) self.assertIsInstance( @@ -96,7 +106,7 @@ class TestDataSource(unittest.TestCase): dataset_args.test_restrict_sequence_id = 0 dataset_args.n_frames_per_sequence = -1 - dataset_args.dataset_root = "manifold://co3d/tree/extracted" + dataset_args.dataset_root = self.dataset_root data_source = ImplicitronDataSource(**args) self.assertIsInstance(