diff --git a/test/lit.cfg.py b/test/lit.cfg.py
index 1caf05784f..defdda42ef 100644
--- a/test/lit.cfg.py
+++ b/test/lit.cfg.py
@@ -6,22 +6,32 @@
 
 import os
 import sys
-import platform
-import re
-import subprocess
-import tempfile
-
+import datetime
+import lit
 import lit.formats
 import lit.util
-
 from lit.llvm import llvm_config
-from lit.llvm.subst import ToolSubst
-from lit.llvm.subst import FindTool
+from lit.formats.base import TestFormat
 
 
 # Configuration file for the 'lit' test runner.
 
+class TimestampedTestFormat(TestFormat):
+    def __init__(self):
+        super().__init__('Timestamped Test')
+
+    def execute(self, test, lit_config):
+        start_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
+        test.setMetric('start_timestamp', start_timestamp)
 
+        # Run the test
+        result = lit.TestRunner.executeShTest(test, lit_config)
+
+        end_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
+        test.setMetric('end_timestamp', end_timestamp)
+
+        return result
+    
 def set_system_desc_features(system_desc):
     config.available_features.add(system_desc["chip_descs"][0]["arch"])
     if len(system_desc["chip_desc_indices"]) > 1:
@@ -31,11 +41,8 @@ def set_system_desc_features(system_desc):
 # name: The name of this test suite.
 config.name = "TTMLIR"
 
-from timestamped_test_format import TimestampedTestFormat
 config.test_format = TimestampedTestFormat()
 
-#config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
-
 # Stablehlo tests can be optionally enabled.
 if config.enable_stablehlo:
     config.available_features.add("stablehlo")
diff --git a/test/timestamped_test_format.py b/test/timestamped_test_format.py
deleted file mode 100644
index b8b16f2a4d..0000000000
--- a/test/timestamped_test_format.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import time
-import datetime
-import lit
-from lit.formats.base import TestFormat
-
-class TimestampedTestFormat(TestFormat):
-    def __init__(self):
-        super().__init__('Timestamped Test')
-
-    def execute(self, test, lit_config):
-        start_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
-        test.setMetric('start_timestamp', start_timestamp)
-
-        # Run the test
-        result = lit.TestRunner.executeShTest(test, lit_config)
-
-        end_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
-        test.setMetric('end_timestamp', end_timestamp)
-
-        return result