From 503880dbbe3b3ceb4700f6377da8a32a4629bd66 Mon Sep 17 00:00:00 2001
From: Liangsheng Yin <hnyls2002@gmail.com>
Date: Thu, 11 Dec 2025 01:00:50 +0900
Subject: [PATCH] [CI] fix UT success check in
 `test_eagle_infer_beta_dp_attention.py` (#14831)

---
 test/nightly/test_eagle_infer_beta_dp_attention_large.py | 9 +--------
 test/srt/test_eagle_infer_beta_dp_attention.py           | 9 +--------
 2 files changed, 2 insertions(+), 16 deletions(-)

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()