code: def assert_node_states(
      state: ClusterResourceState,
      expected_nodes: List[ExpectedNodeState],
  ):
      """
      Assert a GetClusterResourceStateReply has node states that
      matches with the expected nodes.
      """
      assert len(state.node_states) == len(expected_nodes)
  
      # Sort all the nodes by node's node_id
      node_states = sorted(
          state.node_states, key=lambda node: node.node_id
      )
      expected_nodes = sorted(
          expected_nodes, key=lambda node: node.node_id
      )
  
      for actual_node, expected_node in zip(
          node_states, expected_nodes
      ):
          assert actual_node.status == expected_node.node_status
          if expected_node.idle_time_check_cb:
              assert expected_node.idle_time_check_cb(
                  actual_node.idle_duration_ms
              )
  
          if expected_node.labels:
              assert sorted(actual_node.dynamic_labels) == sorted(
                  expected_node.labels
              )
  
imports: import os
  import sys
  import time
  from dataclasses import dataclass
  from typing import Callable, List, Optional, Tuple
  import pytest
  import ray
  import ray._private.ray_constants
  from ray._private.test_utils import wait_for_condition
  from ray.autoscaler.v2.schema import ClusterStatus, LaunchRequest, NodeInfo, ResourceRequestByCount
  from ray.autoscaler.v2.sdk import get_cluster_status, request_cluster_resources
  from ray.autoscaler.v2.tests.util import get_available_resources, get_cluster_resource_state, get_total_resources, report_autoscaling_state
  from ray.core.generated import autoscaler_pb2, autoscaler_pb2_grpc
  from ray.core.generated.autoscaler_pb2 import ClusterResourceState, NodeStatus
  from ray.util.state.api import list_nodes
label_desc: Use image node_id 4 for calling a global function with example usage: assert_node_states(state, expected_nodes) without return types