diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py
index 34ffc5a30..be848c02d 100644
--- a/test/registered/unit/managers/test_hicache_controller_cp.py
+++ b/test/registered/unit/managers/test_hicache_controller_cp.py
@@ -1,6 +1,6 @@
 import sys
 import types
-import unittest
+from unittest import main
 from unittest.mock import patch
 
 import torch
@@ -207,4 +207,4 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    main()