From 0b3a6216f57891d5b34e4d5d1318128829580fc1 Mon Sep 17 00:00:00 2001
From: Jinjing Zhou <VoVAllen@users.noreply.github.com>
Date: Fri, 2 Jul 2021 14:54:38 +0800
Subject: [PATCH] [Test] Enable kvstore test (#3079)

* try enable kvstore test

* fix

* fix

* seperate out kvstore test

* add comment
---
 Jenkinsfile                                                   | 2 +-
 .../{test_new_kvstore.py.bak => kv_store/test_new_kvstore.py} | 0
 tests/distributed/test_dist_graph_store.py                    | 1 +
 tests/scripts/task_unit_test.sh                               | 4 +++-
 4 files changed, 5 insertions(+), 2 deletions(-)
 rename tests/distributed/{test_new_kvstore.py.bak => kv_store/test_new_kvstore.py} (100%)

diff --git a/Jenkinsfile b/Jenkinsfile
index 503e3285ab3d..9d9286a1f5a5 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -57,7 +57,7 @@ def cpp_unit_test_win64() {
 def unit_test_linux(backend, dev) {
   init_git()
   unpack_lib("dgl-${dev}-linux", dgl_linux_libs)
-  timeout(time: 20, unit: 'MINUTES') {
+  timeout(time: 30, unit: 'MINUTES') {
     sh "bash tests/scripts/task_unit_test.sh ${backend} ${dev}"
   }
 }
diff --git a/tests/distributed/test_new_kvstore.py.bak b/tests/distributed/kv_store/test_new_kvstore.py
similarity index 100%
rename from tests/distributed/test_new_kvstore.py.bak
rename to tests/distributed/kv_store/test_new_kvstore.py
diff --git a/tests/distributed/test_dist_graph_store.py b/tests/distributed/test_dist_graph_store.py
index 5137a0c9cbae..99dcf4c36797 100644
--- a/tests/distributed/test_dist_graph_store.py
+++ b/tests/distributed/test_dist_graph_store.py
@@ -176,6 +176,7 @@ def check_dist_emb(g, num_clients, num_nodes, num_edges):
     # Test sparse emb
     try:
         emb = NodeEmbedding(g.number_of_nodes(), 1, 'emb1', emb_init)
+        nids = F.arange(0, int(g.number_of_nodes()))
         lr = 0.001
         optimizer = SparseAdagrad([emb], lr=lr)
         with F.record_grad():
diff --git a/tests/scripts/task_unit_test.sh b/tests/scripts/task_unit_test.sh
index 3b032425f154..45be7c6d1906 100644
--- a/tests/scripts/task_unit_test.sh
+++ b/tests/scripts/task_unit_test.sh
@@ -37,5 +37,7 @@ python3 -m pytest -v --junitxml=pytest_backend.xml tests/$DGLBACKEND || fail "ba
 
 export OMP_NUM_THREADS=1
 if [ $2 != "gpu" ]; then
-    python3 -m pytest -v --capture=tee-sys --junitxml=pytest_distributed.xml tests/distributed || fail "distributed"
+    python3 -m pytest -v --capture=tee-sys --junitxml=pytest_distributed.xml tests/distributed/*.py || fail "distributed"
+    # Seperate kvstore test to another process, to avoid hangs
+    python3 -m pytest -v --capture=tee-sys --junitxml=pytest_distributed.xml tests/distributed/kv_store/*.py || fail "distributed kvstore"
 fi