relink.py: Add --timing option.

Author: samsrabin

relink.py

@@ -9,6 +9,7 @@
 import pwd
 import argparse
 import logging
+import time
 
 DEFAULT_SOURCE_ROOT = "/glade/campaign/cesm/cesmdata/cseg/inputdata/"
 DEFAULT_TARGET_ROOT = (
@@ -185,27 +186,40 @@ def parse_arguments():
         action="store_true",
         help="Show what would be done without making any changes",
     )
+    parser.add_argument(
+        "--timing",
+        action="store_true",
+        help="Measure and display the execution time",
+    )
 
     return parser.parse_args()
 
-
-if __name__ == "__main__":
+def main():
 
     args = parse_arguments()
 
     # Configure logging based on verbosity flags
     if args.quiet:
-        LOG_LEVEL = logging.WARNING
+        log_level = logging.WARNING
     elif args.verbose:
-        LOG_LEVEL = logging.DEBUG
+        log_level = logging.DEBUG
     else:
-        LOG_LEVEL = logging.INFO
+        log_level = logging.INFO
 
-    logging.basicConfig(level=LOG_LEVEL, format="%(message)s", stream=sys.stdout)
+    logging.basicConfig(level=log_level, format="%(message)s", stream=sys.stdout)
 
     my_username = os.environ["USER"]
 
+    start_time = time.time()
+
     # --- Execution ---
     find_and_replace_owned_files(
         args.source_root, args.target_root, my_username, dry_run=args.dry_run
     )
+
+    if args.timing:
+        elapsed_time = time.time() - start_time
+        logger.info("Execution time: %.2f seconds", elapsed_time)
+
+if __name__ == "__main__":
+    main()

tests/test_relink.py

@@ -466,6 +466,94 @@ def test_dry_run_default(self, mock_default_dirs):
             args = relink.parse_arguments()
             assert args.dry_run is False
 
+    def test_timing_flag(self, mock_default_dirs):
+        """Test that --timing flag is parsed correctly."""
+        with patch("sys.argv", ["relink.py", "--timing"]):
+            args = relink.parse_arguments()
+            assert args.timing is True
+
+    def test_timing_default(self, mock_default_dirs):
+        """Test that timing defaults to False."""
+        with patch("sys.argv", ["relink.py"]):
+            args = relink.parse_arguments()
+            assert args.timing is False
+
+
+class TestTiming:
+    """Test suite for timing functionality."""
+
+    @pytest.fixture
+    def temp_dirs(self):
+        """Create temporary source and target directories for testing."""
+        source_dir = tempfile.mkdtemp(prefix="test_source_")
+        target_dir = tempfile.mkdtemp(prefix="test_target_")
+
+        yield source_dir, target_dir
+
+        # Cleanup
+        shutil.rmtree(source_dir, ignore_errors=True)
+        shutil.rmtree(target_dir, ignore_errors=True)
+
+    def test_timing_message_logged(self, tmp_path, caplog):
+        """Test that timing message is logged when timing is enabled."""
+        # Create real directories
+        source_dir = tmp_path / "source"
+        target_dir = tmp_path / "target"
+        source_dir.mkdir()
+        target_dir.mkdir()
+
+        # Create a file
+        source_file = source_dir / "test_file.txt"
+        target_file = target_dir / "test_file.txt"
+        source_file.write_text("source")
+        target_file.write_text("target")
+
+        # Mock sys.argv to simulate running with --timing
+        test_argv = [
+            "relink.py",
+            "--source-root", str(source_dir),
+            "--target-root", str(target_dir),
+            "--timing"
+        ]
+
+        with patch("sys.argv", test_argv):
+            with caplog.at_level(logging.INFO):
+                # Call main() which includes the timing logic
+                relink.main()
+
+        # Verify timing message was logged
+        assert "Execution time:" in caplog.text
+        assert "seconds" in caplog.text
+
+    def test_timing_not_logged_by_default(self, tmp_path, caplog):
+        """Test that timing message is not logged when timing is disabled."""
+        # Create real directories
+        source_dir = tmp_path / "source"
+        target_dir = tmp_path / "target"
+        source_dir.mkdir()
+        target_dir.mkdir()
+
+        # Create a file
+        source_file = source_dir / "test_file.txt"
+        target_file = target_dir / "test_file.txt"
+        source_file.write_text("source")
+        target_file.write_text("target")
+
+        # Mock sys.argv WITHOUT --timing flag
+        test_argv = [
+            "relink.py",
+            "--source-root", str(source_dir),
+            "--target-root", str(target_dir)
+        ]
+
+        with patch("sys.argv", test_argv):
+            with caplog.at_level(logging.INFO):
+                # Call main() without timing flag
+                relink.main()
+
+        # Verify timing message was NOT logged
+        assert "Execution time:" not in caplog.text
+
 
 class TestDryRun:
     """Test suite for dry-run functionality."""