diff --git a/test/nightly/test_eagle_infer_beta_dp_attention_large.py b/test/nightly/test_eagle_infer_beta_dp_attention_large.py
index 00539cc01..6bb02e38f 100644
--- a/test/nightly/test_eagle_infer_beta_dp_attention_large.py
+++ b/test/nightly/test_eagle_infer_beta_dp_attention_large.py
@@ -98,11 +98,4 @@ class TestEagleDPAttnServerLarge(CustomTestCase):
 
 
 if __name__ == "__main__":
-    s = unittest.TestSuite()
-    large_test = unittest.defaultTestLoader.loadTestsFromTestCase(
-        TestEagleDPAttnServerLarge
-    )
-    s.addTest(large_test)
-
-    runner = unittest.TextTestRunner()
-    runner.run(s)
+    unittest.main()
diff --git a/test/srt/test_eagle_infer_beta_dp_attention.py b/test/srt/test_eagle_infer_beta_dp_attention.py
index 96ecdb2f7..038f683fe 100644
--- a/test/srt/test_eagle_infer_beta_dp_attention.py
+++ b/test/srt/test_eagle_infer_beta_dp_attention.py
@@ -80,11 +80,4 @@ class TestEagleDPAttnServerSmall(CustomTestCase):
 
 
 if __name__ == "__main__":
-    s = unittest.TestSuite()
-    small_test = unittest.defaultTestLoader.loadTestsFromTestCase(
-        TestEagleDPAttnServerSmall
-    )
-    s.addTest(small_test)
-
-    runner = unittest.TextTestRunner()
-    runner.run(s)
+    unittest.main()