github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/tests/unit/sdk/test_cluster.py (about) 1 import unittest 2 from typing import List, Optional 3 from unittest.mock import Mock, patch 4 5 from aistore.sdk.bucket import Bucket 6 from aistore.sdk.cluster import Cluster 7 from aistore.sdk.const import ( 8 HTTP_METHOD_GET, 9 QPARAM_WHAT, 10 QPARAM_PROVIDER, 11 ACT_LIST, 12 PROVIDER_AIS, 13 WHAT_SMAP, 14 URL_PATH_DAEMON, 15 URL_PATH_BUCKETS, 16 URL_PATH_HEALTH, 17 QPARAM_PRIMARY_READY_REB, 18 URL_PATH_CLUSTER, 19 WHAT_ALL_XACT_STATUS, 20 WHAT_ALL_RUNNING_STATUS, 21 URL_PATH_ETL, 22 ) 23 from aistore.sdk.request_client import RequestClient 24 from aistore.sdk.types import ( 25 NodeStats, 26 NodeTracker, 27 NodeCapacity, 28 Smap, 29 ActionMsg, 30 BucketModel, 31 JobStatus, 32 JobQuery, 33 ETLInfo, 34 Snode, 35 NetInfo, 36 NodeCounter, 37 NodeLatency, 38 NodeThroughput, 39 ) 40 41 from tests.utils import test_cases 42 43 44 class TestCluster(unittest.TestCase): # pylint: disable=unused-variable 45 def setUp(self) -> None: 46 self.mock_client = Mock(RequestClient) 47 self.cluster = Cluster(self.mock_client) 48 49 def test_get_info(self): 50 expected_result = Mock() 51 self.mock_client.request_deserialize.return_value = expected_result 52 result = self.cluster.get_info() 53 self.assertEqual(result, expected_result) 54 self.mock_client.request_deserialize.assert_called_with( 55 HTTP_METHOD_GET, 56 path=URL_PATH_DAEMON, 57 res_model=Smap, 58 params={QPARAM_WHAT: WHAT_SMAP}, 59 ) 60 61 def test_list_buckets(self): 62 provider = "any-provider" 63 expected_params = {QPARAM_PROVIDER: provider} 64 self.list_buckets_exec_assert(expected_params, provider=provider) 65 66 def test_list_buckets_default_param(self): 67 expected_params = {QPARAM_PROVIDER: PROVIDER_AIS} 68 self.list_buckets_exec_assert(expected_params) 69 70 def list_buckets_exec_assert(self, expected_params, **kwargs): 71 expected_result = [Mock(Bucket)] 72 self.mock_client.request_deserialize.return_value = expected_result 73 74 res = self.cluster.list_buckets(**kwargs) 75 76 self.assertEqual(expected_result, res) 77 self.mock_client.request_deserialize.assert_called_with( 78 HTTP_METHOD_GET, 79 path=URL_PATH_BUCKETS, 80 res_model=List[BucketModel], 81 json=ActionMsg(action=ACT_LIST).dict(), 82 params=expected_params, 83 ) 84 85 def test_is_ready_exception(self): 86 self.mock_client.request.side_effect = Exception 87 self.assertFalse(self.cluster.is_ready()) 88 89 @test_cases(True, False) 90 def test_is_ready(self, test_case): 91 expected_params = {QPARAM_PRIMARY_READY_REB: "true"} 92 primary_proxy_endpoint = "primary_proxy_url" 93 94 mock_response = Mock() 95 mock_response.ok = test_case 96 self.mock_client.request.return_value = mock_response 97 mock_smap = Mock(spec=Smap) 98 mock_snode = Mock(spec=Snode) 99 mock_netinfo = Mock(spec=NetInfo) 100 mock_netinfo.direct_url = primary_proxy_endpoint 101 mock_snode.public_net = mock_netinfo 102 mock_smap.proxy_si = mock_snode 103 self.mock_client.request_deserialize.return_value = mock_smap 104 105 self.assertEqual(test_case, self.cluster.is_ready()) 106 self.mock_client.request.assert_called_with( 107 HTTP_METHOD_GET, 108 path=URL_PATH_HEALTH, 109 endpoint=primary_proxy_endpoint, 110 params=expected_params, 111 ) 112 113 def test_list_jobs_status_default_params(self): 114 expected_request_val = JobQuery().as_dict() 115 self._list_jobs_status_exec_assert(expected_request_val) 116 117 def test_list_jobs_status(self): 118 job_kind = "kind" 119 target_id = "specific_node" 120 121 expected_request_val = JobQuery(kind=job_kind, target=target_id).as_dict() 122 self._list_jobs_status_exec_assert( 123 expected_request_val, 124 job_kind=job_kind, 125 target_id=target_id, 126 ) 127 128 def test_list_jobs_status_no_result(self): 129 self.mock_client.request_deserialize.return_value = None 130 self.assertEqual([], self.cluster.list_jobs_status()) 131 132 def _list_jobs_status_exec_assert(self, expected_request_val, **kwargs): 133 returned_status = JobStatus() 134 self.mock_client.request_deserialize.return_value = returned_status 135 136 res = self.cluster.list_jobs_status(**kwargs) 137 138 self.assertEqual(returned_status, res) 139 self.mock_client.request_deserialize.assert_called_with( 140 HTTP_METHOD_GET, 141 path=URL_PATH_CLUSTER, 142 res_model=Optional[List[JobStatus]], 143 json=expected_request_val, 144 params={QPARAM_WHAT: WHAT_ALL_XACT_STATUS}, 145 ) 146 147 def test_list_running_jobs_default_params(self): 148 expected_request_val = JobQuery(active=True).as_dict() 149 self._list_running_jobs_exec_assert(expected_request_val) 150 151 def test_list_running_jobs(self): 152 job_kind = "job-kind" 153 target_id = "my-target" 154 expected_request_val = JobQuery( 155 active=True, kind=job_kind, target=target_id 156 ).as_dict() 157 self._list_running_jobs_exec_assert( 158 expected_request_val, job_kind=job_kind, target_id=target_id 159 ) 160 161 def _list_running_jobs_exec_assert(self, expected_request_val, **kwargs): 162 mock_response = ["job_1_kind[job_1_id]", "job_2_kind[job_2_id]"] 163 self.mock_client.request_deserialize.return_value = mock_response 164 165 res = self.cluster.list_running_jobs(**kwargs) 166 167 self.assertEqual(mock_response, res) 168 self.mock_client.request_deserialize.assert_called_with( 169 HTTP_METHOD_GET, 170 path=URL_PATH_CLUSTER, 171 res_model=List[str], 172 json=expected_request_val, 173 params={QPARAM_WHAT: WHAT_ALL_RUNNING_STATUS}, 174 ) 175 176 def test_list_running_etls(self): 177 mock_response = Mock() 178 self.mock_client.request_deserialize.return_value = mock_response 179 response = self.cluster.list_running_etls() 180 self.assertEqual(mock_response, response) 181 self.mock_client.request_deserialize.assert_called_with( 182 HTTP_METHOD_GET, path=URL_PATH_ETL, res_model=List[ETLInfo] 183 ) 184 185 @patch("aistore.sdk.cluster.Cluster._get_smap") 186 def test_get_performance(self, mock_get_smap): 187 mock_targets = ["target1", "target2"] 188 mock_smap = Smap( 189 tmap={"target1": Mock(spec=Snode), "target2": Mock(spec=Snode)}, 190 pmap={"proxy1": Mock(spec=Snode)}, 191 proxy_si=Mock(spec=Snode), 192 ) 193 mock_get_smap.return_value = mock_smap 194 195 mock_node_tracker = NodeTracker( 196 append_ns=1000, 197 del_n=10, 198 disk_sdb_util=50.0, 199 disk_sdb_avg_rsize=1024, 200 disk_sdb_avg_wsize=2048, 201 disk_sdb_read_bps=1000000, 202 disk_sdb_write_bps=500000, 203 dl_ns=2000, 204 dsort_creation_req_n=5, 205 dsort_creation_resp_n=5, 206 dsort_creation_resp_ns=100, 207 dsort_extract_shard_mem_n=2, 208 dsort_extract_shard_size=102400, 209 err_del_n=1, 210 err_get_n=2, 211 get_bps=1024000, 212 get_cold_n=20, 213 get_cold_rw_ns=3000, 214 get_cold_size=204800, 215 get_n=100, 216 get_ns=4000, 217 get_redir_ns=500, 218 get_size=409600, 219 kalive_ns=600, 220 lcache_evicted_n=15, 221 lcache_flush_cold_n=10, 222 lru_evict_n=5, 223 lru_evict_size=102400, 224 lst_n=50, 225 lst_ns=700, 226 put_bps=2048000, 227 put_n=80, 228 put_ns=8000, 229 put_redir_ns=600, 230 put_size=819200, 231 remote_deleted_del_n=3, 232 stream_in_n=40, 233 stream_in_size=409600, 234 stream_out_n=35, 235 stream_out_size=204800, 236 up_ns_time=10000, 237 ver_change_n=25, 238 ver_change_size=512000, 239 ) 240 241 mock_node_stats = NodeStats( 242 snode=Mock(spec=Snode), 243 tracker=mock_node_tracker, 244 capacity=Mock(spec=NodeCapacity), 245 rebalance_snap={}, 246 status="", 247 deployment="", 248 ais_version="", 249 build_time="", 250 k8s_pod_name="", 251 sys_info={}, 252 smap_version="", 253 ) 254 255 self.mock_client.request_deserialize.return_value = mock_node_stats 256 257 performance = self.cluster.get_performance() 258 259 for target_id in mock_targets: 260 throughput = performance.throughput[target_id].as_dict() 261 latency = performance.latency[target_id].as_dict() 262 counter = performance.counters[target_id].as_dict() 263 264 self.assertEqual(NodeThroughput(mock_node_tracker).as_dict(), throughput) 265 self.assertEqual(NodeLatency(mock_node_tracker).as_dict(), latency) 266 self.assertEqual(NodeCounter(mock_node_tracker).as_dict(), counter) 267 268 self.mock_client.request_deserialize.assert_called()