Add regression tests for get_latest and get_oldest functions

fms_fsdp/utils/checkpointing_utils.py

@@ -25,12 +25,8 @@ def get_latest(targdir, qualifier=lambda x: True):
     If directory is empty or nonexistent or no items qualify, return None."""
     if os.path.exists(targdir) and len(os.listdir(targdir)) > 0:
         latest = max(
-            [
-                os.path.join(targdir, x)
-                for x in os.listdir(targdir)
-                if qualifier(os.path.join(targdir, x))
-            ],
-            key=lambda path: int(path.split("/")[-1].split("_")[1]),
+            [x for x in os.listdir(targdir) if qualifier(os.path.join(targdir, x))],
+            key=lambda path: int(path.split("_")[1]),
         )
         return os.path.join(targdir, latest)
     return None
@@ -41,11 +37,7 @@ def get_oldest(targdir, qualifier=lambda x: True):
     If directory is empty or nonexistent or no items qualify, return None."""
     if os.path.exists(targdir) and len(os.listdir(targdir)) > 0:
         oldest = min(
-            [
-                os.path.join(targdir, x)
-                for x in os.listdir(targdir)
-                if qualifier(os.path.join(targdir, x))
-            ],
+            [x for x in os.listdir(targdir) if qualifier(os.path.join(targdir, x))],
             key=os.path.getctime,
         )
         return os.path.join(targdir, oldest)

tests/test_utils_checkpointing.py

@@ -0,0 +1,36 @@
+"""
+Tests for functions in utils/checkpointing_utils.py
+"""
+
+import os
+import tempfile
+
+from fms_fsdp.utils.checkpointing_utils import get_latest, get_oldest
+
+
+def test_get_oldest():
+    """
+    Ensure that the get_oldest function returns the name of the file with the oldest
+    timestamp (i.e. that was created first).
+    """
+    with tempfile.TemporaryDirectory() as tempdir:
+        for i in range(3):
+            filename = os.path.join(tempdir, f"file_{i}")
+            print("random content", file=open(file=filename, mode="w"))
+
+        oldest_filename = get_oldest(targdir=tempdir)
+        assert oldest_filename.endswith("file_0")
+
+
+def test_get_latest():
+    """
+    Ensure that the get_latest function returns the name of the file with the latest
+    integer suffix (i.e. that was created last).
+    """
+    with tempfile.TemporaryDirectory() as tempdir:
+        for i in range(3):
+            filename = os.path.join(tempdir, f"file_{i}")
+            print("random content", file=open(file=filename, mode="w"))
+
+        latest_filename = get_latest(targdir=tempdir)
+        assert latest_filename.endswith("file_2")