diff --git a/tests/cpp/common/test_quantile.h b/tests/cpp/common/test_quantile.h
index e91f19ef8..083766cfb 100644
--- a/tests/cpp/common/test_quantile.h
+++ b/tests/cpp/common/test_quantile.h
@@ -16,10 +16,18 @@ inline void InitRabitContext(std::string msg, int32_t n_workers) {
     LOG(WARNING) << msg << " as `DMLC_TRACKER_PORT` is not set up.";
     return;
   }
+  auto uri = std::getenv("DMLC_TRACKER_URI");
+  std::string uri_str;
+  if (uri) {
+    uri_str = uri;
+  } else {
+    LOG(WARNING) << msg << " as `DMLC_TRACKER_URI` is not set up.";
+    return;
+  }
 
   std::vector<std::string> envs{
       "DMLC_TRACKER_PORT=" + port_str,
-      "DMLC_TRACKER_URI=127.0.0.1",
+      "DMLC_TRACKER_URI=" + uri_str,
       "DMLC_NUM_WORKER=" + std::to_string(n_workers)};
   char* c_envs[] {&(envs[0][0]), &(envs[1][0]), &(envs[2][0])};
   rabit::Init(3, c_envs);
diff --git a/tests/python/test_ranking.py b/tests/python/test_ranking.py
index f68c47c3f..98bca122f 100644
--- a/tests/python/test_ranking.py
+++ b/tests/python/test_ranking.py
@@ -136,9 +136,9 @@ class TestRanking:
         cv = xgboost.cv(self.params, self.dtrain, num_boost_round=2500,
                         early_stopping_rounds=10, nfold=10, as_pandas=False)
         assert isinstance(cv, dict)
-        assert (set(cv.keys()) == {'test-ndcg-mean', 'train-ndcg-mean', 'test-ndcg-std',
-                                   'train-ndcg-std'},
-                'CV results dict key mismatch.')
+        assert set(cv.keys()) == {
+            'test-ndcg-mean', 'train-ndcg-mean', 'test-ndcg-std', 'train-ndcg-std'
+        }, "CV results dict key mismatch."
 
     def test_cv_no_shuffle(self):
         """