Add test cases to test query_running_nodes

dlrover/python/tests/test_servicer.py

@@ -38,10 +38,14 @@ class MasterServicerTest(unittest.TestCase):
  def setUp(self) -> None:
  mock_k8s_client()
  params = MockK8sPSJobArgs()
+ worker_resource = params.node_args[NodeType.WORKER].group_resource
+ worker_resource.node_resource.gpu_num = 1
+ worker_resource.node_resource.gpu_type = "a100"
  params.initilize()
  speed_monitor = SpeedMonitor()
  self.task_manager = TaskManager(False, speed_monitor)
  self.job_manager = create_job_manager(params, speed_monitor)
+ self.job_manager._init_nodes()
  self.job_metric_collector = JobMetricCollector(
  "1", "default", "local", "dlrover"
  )
@@ -55,6 +59,13 @@ def setUp(self) -> None:
  elastic_ps_service=self.elastic_ps_service,
  )
 
+ def test_query_running_nodes(self):
+ req = empty_pb2.Empty()
+ workers = self.job_manager._job_nodes[NodeType.WORKER]
+ workers[0].status = NodeStatus.RUNNING
+ res = self.servicer.query_running_nodes(req, None)
+ self.assertEqual(len(res.nodes), 0)
+
  def test_dataset_service(self):
  request = elastic_training_pb2.ReportDatasetShardParamsRequest()
  request.batch_size = 10