Fixed tests.

mmcdermott · Aug 29, 2024 · 5f38154 · 5f38154
1 parent 24b8eeb
commit 5f38154
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/src/MEDS_transforms/mapreduce/utils.py b/src/MEDS_transforms/mapreduce/utils.py
@@ -453,10 +453,11 @@ def shard_iterator(
         >>> includes_only_train
         False
 
-    If it can't find any files, it will return an empty list:
+    If it can't find any files, it will error:
         >>> fps, includes_only_train = shard_iterator(cfg)
-        >>> fps
-        []
+        Traceback (most recent call last):
+            ...
+        FileNotFoundError: No shards found in ... with suffix .parquet. Directory contents:...
     """
 
     input_dir = Path(cfg.stage_cfg.data_input_dir)

diff --git a/src/MEDS_transforms/runner.py b/src/MEDS_transforms/runner.py
@@ -242,8 +242,8 @@ def fix_str_for_path(s: str) -> str:
     return s.replace(" ", "_").replace("/", ".")
 
 
-if __name__ == "__main__":
-    OmegaConf.register_new_resolver("load_yaml_file", load_yaml_file, replace=False)
-    OmegaConf.register_new_resolver("fix_str_for_path", fix_str_for_path, replace=False)
+OmegaConf.register_new_resolver("load_yaml_file", load_yaml_file, replace=False)
+OmegaConf.register_new_resolver("fix_str_for_path", fix_str_for_path, replace=False)
 
+if __name__ == "__main__":
     main()