Add tests

Author: misrasaurabh1

codeflash/code_utils/checkpoint.py

@@ -11,7 +11,7 @@
 
 
 class CodeflashRunCheckpoint:
-    def __init__(self, module_root: Path, checkpoint_dir: str = "/tmp") -> None:
+    def __init__(self, module_root: Path, checkpoint_dir: Path = Path("/tmp")) -> None:
         self.module_root = module_root
         self.checkpoint_dir = Path(checkpoint_dir)
         # Create a unique checkpoint file name

tests/test_codeflash_checkpoint.py

@@ -0,0 +1,176 @@
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint, get_all_historical_functions
+
+
+class TestCodeflashRunCheckpoint:
+    @pytest.fixture
+    def temp_dir(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            yield Path(temp_dir)
+
+    def test_initialization(self, temp_dir):
+        module_root = Path("/fake/module/root")
+        checkpoint = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir)
+
+        # Check if checkpoint file was created
+        assert checkpoint.checkpoint_path.exists()
+
+        # Check if metadata was written correctly
+        with open(checkpoint.checkpoint_path) as f:
+            metadata = json.loads(f.readline())
+            assert metadata["type"] == "metadata"
+            assert metadata["module_root"] == str(module_root)
+            assert "created_at" in metadata
+            assert "last_updated" in metadata
+
+    def test_add_function_to_checkpoint(self, temp_dir):
+        module_root = Path("/fake/module/root")
+        checkpoint = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir)
+
+        # Add a function to the checkpoint
+        function_name = "module.submodule.function"
+        checkpoint.add_function_to_checkpoint(function_name, status="optimized")
+
+        # Read the checkpoint file and verify
+        with open(checkpoint.checkpoint_path) as f:
+            lines = f.readlines()
+            assert len(lines) == 2  # Metadata + function entry
+
+            function_data = json.loads(lines[1])
+            assert function_data["type"] == "function"
+            assert function_data["function_name"] == function_name
+            assert function_data["status"] == "optimized"
+            assert "timestamp" in function_data
+
+    def test_add_function_with_additional_info(self, temp_dir):
+        module_root = Path("/fake/module/root")
+        checkpoint = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir)
+
+        # Add a function with additional info
+        function_name = "module.submodule.function"
+        additional_info = {"execution_time": 1.5, "memory_usage": "10MB"}
+        checkpoint.add_function_to_checkpoint(function_name, status="optimized", additional_info=additional_info)
+
+        # Read the checkpoint file and verify
+        with open(checkpoint.checkpoint_path) as f:
+            lines = f.readlines()
+            function_data = json.loads(lines[1])
+            assert function_data["execution_time"] == 1.5
+            assert function_data["memory_usage"] == "10MB"
+
+    def test_update_metadata_timestamp(self, temp_dir):
+        module_root = Path("/fake/module/root")
+        checkpoint = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir)
+
+        # Get initial timestamp
+        with open(checkpoint.checkpoint_path) as f:
+            initial_metadata = json.loads(f.readline())
+            initial_timestamp = initial_metadata["last_updated"]
+
+        # Wait a bit to ensure timestamp changes
+        import time
+
+        time.sleep(0.01)
+
+        # Update timestamp
+        checkpoint._update_metadata_timestamp()
+
+        # Check if timestamp was updated
+        with open(checkpoint.checkpoint_path) as f:
+            updated_metadata = json.loads(f.readline())
+            updated_timestamp = updated_metadata["last_updated"]
+
+        assert updated_timestamp > initial_timestamp
+
+    def test_cleanup(self, temp_dir):
+        module_root = Path("/fake/module/root")
+
+        # Create multiple checkpoint files
+        checkpoint1 = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir)
+        checkpoint2 = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir)
+
+        # Create a checkpoint for a different module
+        different_module = Path("/different/module")
+        checkpoint3 = CodeflashRunCheckpoint(different_module, checkpoint_dir=temp_dir)
+
+        # Verify all files exist
+        assert checkpoint1.checkpoint_path.exists()
+        assert checkpoint2.checkpoint_path.exists()
+        assert checkpoint3.checkpoint_path.exists()
+
+        # Clean up files for module_root
+        checkpoint1.cleanup()
+
+        # Check that only the files for module_root were deleted
+        assert not checkpoint1.checkpoint_path.exists()
+        assert not checkpoint2.checkpoint_path.exists()
+        assert checkpoint3.checkpoint_path.exists()
+
+
+class TestGetAllHistoricalFunctions:
+    @pytest.fixture
+    def setup_checkpoint_files(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_dir_path = Path(temp_dir)
+            module_root = Path("/fake/module/root")
+
+            # Create a checkpoint file with some functions
+            checkpoint = CodeflashRunCheckpoint(module_root, checkpoint_dir=temp_dir_path)
+            checkpoint.add_function_to_checkpoint("module.func1", status="optimized")
+            checkpoint.add_function_to_checkpoint("module.func2", status="failed")
+
+            # Create an old checkpoint file (more than 7 days old)
+            old_checkpoint_path = temp_dir_path / "codeflash_checkpoint_old.jsonl"
+            with open(old_checkpoint_path, "w") as f:
+                # Create metadata with old timestamp (8 days ago)
+                import time
+
+                old_time = time.time() - (8 * 24 * 60 * 60)
+                metadata = {
+                    "type": "metadata",
+                    "module_root": str(module_root),
+                    "created_at": old_time,
+                    "last_updated": old_time,
+                }
+                f.write(json.dumps(metadata) + "\n")
+
+                # Add a function entry
+                function_data = {
+                    "type": "function",
+                    "function_name": "module.old_func",
+                    "status": "optimized",
+                    "timestamp": old_time,
+                }
+                f.write(json.dumps(function_data) + "\n")
+
+            # Create a checkpoint for a different module
+            different_module = Path("/different/module")
+            diff_checkpoint = CodeflashRunCheckpoint(different_module, checkpoint_dir=temp_dir_path)
+            diff_checkpoint.add_function_to_checkpoint("different.func", status="optimized")
+
+            yield module_root, temp_dir_path
+
+    def test_get_all_historical_functions(self, setup_checkpoint_files):
+        module_root, checkpoint_dir = setup_checkpoint_files
+
+        # Get historical functions
+        functions = get_all_historical_functions(module_root, checkpoint_dir)
+
+        # Verify the functions from the current checkpoint are included
+        assert "module.func1" in functions
+        assert "module.func2" in functions
+        assert functions["module.func1"]["status"] == "optimized"
+        assert functions["module.func2"]["status"] == "failed"
+
+        # Verify the old function is not included (file should be deleted)
+        assert "module.old_func" not in functions
+
+        # Verify the function from the different module is not included
+        assert "different.func" not in functions
+
+        # Verify the old checkpoint file was deleted
+        assert not (checkpoint_dir / "codeflash_checkpoint_old.jsonl").exists()