69 lines
2.4 KiB
Python
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()
|