Files
2026-07-13 13:17:40 +08:00

69 lines
2.4 KiB
Python

import unittest
from unittest.mock import Mock
from ray.autoscaler._private.util import get_per_node_breakdown_as_dict
class TestGetPerNodeBreakdown(unittest.TestCase):
def setUp(self):
# Create a mock LoadMetricsSummary object with the required attributes
lm_summary_mock_data = {
"e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": {
"node:172.31.6.57": [0.0, 1.0],
"object_store_memory": [0.0, 13984228147.0],
"memory": [0.0, 27968456295.0],
"node:__internal_head__": [0.0, 1.0],
"CPU": [0.0, 8.0],
}
}
self.lm_summary_mock = Mock()
self.lm_summary_mock.usage_by_node = lm_summary_mock_data
def test_get_per_node_breakdown_as_dict(self):
result = get_per_node_breakdown_as_dict(self.lm_summary_mock)
expected_output = {
"e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": (
"0.0/8.0 CPU\n0B/26.05GiB memory\n0B/13.02GiB object_store_memory"
)
}
self.assertEqual(result, expected_output)
def test_get_per_node_breakdown_as_dict_empty_summary(self):
# Test with an empty lm_summary
lm_summary_mock_data = {}
self.lm_summary_mock.usage_by_node = lm_summary_mock_data
result = get_per_node_breakdown_as_dict(self.lm_summary_mock)
expected_output = {}
self.assertEqual(result, expected_output)
def test_get_per_node_breakdown_as_dict_missing_usage(self):
# Test with missing usage data for a node
lm_summary_mock_data = {
"e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": {
"node:172.31.6.57": [0.0, 1.0],
"object_store_memory": [0.0, 13984228147.0],
# 'memory': [0.0, 27968456295.0], # Missing memory data
"node:__internal_head__": [0.0, 1.0],
"CPU": [0.0, 8.0],
}
}
self.lm_summary_mock.usage_by_node = lm_summary_mock_data
result = get_per_node_breakdown_as_dict(self.lm_summary_mock)
expected_output = {
"e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": "0.0/8.0 CPU\n"
"0B/13.02GiB object_store_memory"
}
self.assertEqual(result, expected_output)
if __name__ == "__main__":
unittest.main()