diff --git a/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt b/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt
index ce130243fefae98d2c4e3c7f3710cecd962509eb..fb868e82ce7cfbd6c50de4be7cf54f1d80d3117e 100755
--- a/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt
+++ b/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt
@@ -4,12 +4,16 @@ file(
   "test_*.py")
 string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
 
+if(NOT WITH_HETERPS)
+  list(REMOVE_ITEM TEST_OPS "test_gpubox_ps")
+endif()
+
 foreach(TEST_OP ${TEST_OPS})
   py_test_modules(${TEST_OP} MODULES ${TEST_OP})
   list(APPEND TEST_OPS ${TEST_OP})
   set_tests_properties(${TEST_OP} PROPERTIES TIMEOUT 50)
 endforeach()
 
-if(WITH_PSCORE)
+if(WITH_HETERPS)
   set_tests_properties(test_gpubox_ps PROPERTIES LABELS "RUN_TYPE=GPUPS")
 endif()
diff --git a/python/paddle/fluid/tests/unittests/ps/gpubox_run.sh b/python/paddle/fluid/tests/unittests/ps/gpubox_run.sh
index d3cfcf38aa4c00e677967c7cfed7495083e1874b..55a49c5cd8a632e1f5ad117b56928b090a2e1a33 100644
--- a/python/paddle/fluid/tests/unittests/ps/gpubox_run.sh
+++ b/python/paddle/fluid/tests/unittests/ps/gpubox_run.sh
@@ -31,8 +31,8 @@ export PADDLE_PSERVERS_IP_PORT_LIST="127.0.0.1:29011"
 export PADDLE_PSERVER_PORT_ARRAY=(29011)
 
 # set gpu numbers according to your device
-export FLAGS_selected_gpus="0,1,2,3,4,5,6,7"
-#export FLAGS_selected_gpus="0,1"
+#export FLAGS_selected_gpus="0,1,2,3,4,5,6,7"
+export FLAGS_selected_gpus="0,1"
 
 # set your model yaml
 #SC="gpubox_ps_trainer.py"
@@ -57,4 +57,8 @@ do
     python3.7 -u $SC &> ./log/worker.$i.log
 done
 
+if [ $? -eq 0 ];then
 echo "Training log stored in ./log/"
+else
+exit 1
+fi
diff --git a/python/paddle/fluid/tests/unittests/ps/test_gpubox_ps.py b/python/paddle/fluid/tests/unittests/ps/test_gpubox_ps.py
index 8ebc6ba8a342414c8c123e2069ef68ca029de469..7c7fedd812c6d12c75d23fab22fbbc5763574ef8 100755
--- a/python/paddle/fluid/tests/unittests/ps/test_gpubox_ps.py
+++ b/python/paddle/fluid/tests/unittests/ps/test_gpubox_ps.py
@@ -21,15 +21,16 @@ import unittest
 
 class GpuBoxTest(unittest.TestCase):
     def test_gpubox(self):
-        exitcode = os.system('sh gpubox_run.sh')
-        os.system('rm *_train_desc.prototxt')
+        if not os.path.exists('./train_data'):
+            os.system('bash download_criteo_data.sh')
+
+        exitcode = os.system('bash gpubox_run.sh')
         if os.path.exists('./train_data'):
             os.system('rm -rf train_data')
         if os.path.exists('./log'):
             os.system('rm -rf log')
+        assert exitcode == 0
 
 
 if __name__ == '__main__':
-    if not os.path.exists('./train_data'):
-        os.system('sh download_criteo_data.sh')
     unittest.main()