github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/tests/unit/sdk/test_utils.py (about) 1 import json 2 import unittest 3 from unittest.mock import Mock, patch, mock_open 4 5 from msgspec import msgpack 6 from requests import Response 7 8 from aistore.sdk import utils 9 from aistore.sdk.const import MSGPACK_CONTENT_TYPE, HEADER_CONTENT_TYPE 10 from aistore.sdk.errors import ( 11 AISError, 12 ErrRemoteBckNotFound, 13 ErrBckNotFound, 14 ErrBckAlreadyExists, 15 ErrETLAlreadyExists, 16 ) 17 from tests.const import PREFIX_NAME 18 from tests.utils import test_cases 19 20 21 # pylint: disable=unused-variable 22 class TestUtils(unittest.TestCase): 23 def test_handle_error_no_text(self): 24 mock_response = Mock(text="") 25 26 utils.handle_errors(mock_response) 27 mock_response.raise_for_status.assert_called() 28 29 def test_handle_error_decode_err(self): 30 err_status = 300 31 err_msg = "error message iso-8859-1" 32 expected_text = json.dumps({"status": err_status, "message": err_msg}) 33 # Fail initial decoding, then return the decoded text 34 decode_err = UnicodeDecodeError("1", b"2", 3, 4, "5") 35 mock_iso_text = Mock(spec=bytes) 36 mock_iso_text.decode.side_effect = [decode_err, expected_text] 37 self.handle_err_exec_assert(AISError, err_status, err_msg, mock_iso_text) 38 39 @test_cases(399, 500) 40 def test_handle_error_ais_err(self, err_status): 41 err_msg = "error message" 42 expected_text = json.dumps({"status": err_status, "message": err_msg}) 43 mock_text = Mock(spec=bytes) 44 mock_text.decode.return_value = expected_text 45 self.handle_err_exec_assert(AISError, err_status, err_msg, mock_text) 46 47 @test_cases( 48 ("cloud bucket does not exist", ErrRemoteBckNotFound), 49 ("remote bucket does not exist", ErrRemoteBckNotFound), 50 ("bucket does not exist", ErrBckNotFound), 51 ("bucket already exists", ErrBckAlreadyExists), 52 ("etl already exists", ErrETLAlreadyExists), 53 ) 54 def test_handle_error_no_remote_bucket(self, test_case): 55 err_msg, expected_err = test_case 56 err_status = 400 57 expected_text = json.dumps({"status": err_status, "message": err_msg}) 58 mock_text = Mock(spec=bytes) 59 mock_text.decode.return_value = expected_text 60 self.handle_err_exec_assert(expected_err, err_status, err_msg, mock_text) 61 62 def handle_err_exec_assert(self, err_type, err_status, err_msg, mock_err_text): 63 mock_response = Mock(text=mock_err_text) 64 with self.assertRaises(err_type) as context: 65 utils.handle_errors(mock_response) 66 self.assertEqual(err_msg, context.exception.message) 67 self.assertEqual(err_status, context.exception.status_code) 68 69 @test_cases((0, 0.1), (-1, 0.1), (64, 1), (128, 2), (100000, 1562.5)) 70 def test_probing_frequency(self, test_case): 71 self.assertEqual(test_case[1], utils.probing_frequency(test_case[0])) 72 73 @patch("pathlib.Path.is_file") 74 @patch("pathlib.Path.exists") 75 def test_validate_file(self, mock_exists, mock_is_file): 76 mock_exists.return_value = False 77 with self.assertRaises(ValueError): 78 utils.validate_file("any path") 79 mock_exists.return_value = True 80 mock_is_file.return_value = False 81 with self.assertRaises(ValueError): 82 utils.validate_file("any path") 83 mock_is_file.return_value = True 84 utils.validate_file("any path") 85 86 @patch("pathlib.Path.is_dir") 87 @patch("pathlib.Path.exists") 88 def test_validate_dir(self, mock_exists, mock_is_dir): 89 mock_exists.return_value = False 90 with self.assertRaises(ValueError): 91 utils.validate_directory("any path") 92 mock_exists.return_value = True 93 mock_is_dir.return_value = False 94 with self.assertRaises(ValueError): 95 utils.validate_directory("any path") 96 mock_is_dir.return_value = True 97 utils.validate_directory("any path") 98 99 def test_read_file_bytes(self): 100 data = b"Test data" 101 with patch("builtins.open", mock_open(read_data=data)): 102 res = utils.read_file_bytes("any path") 103 self.assertEqual(data, res) 104 105 @test_cases((123, "123 Bytes"), (None, "unknown")) 106 def test_get_file_size(self, test_case): 107 mock_file = Mock() 108 mock_file.stat.return_value = Mock(st_size=test_case[0]) 109 self.assertEqual(test_case[1], utils.get_file_size(mock_file)) 110 111 @test_cases( 112 (PREFIX_NAME, [PREFIX_NAME], None), 113 ("prefix-{}", ["prefix-{}"], None), 114 ("prefix-{0..1..2..3}", ["prefix-{0..1..2..3}"], None), 115 ("prefix-{0..1..2}}", [], ValueError), 116 ( 117 "prefix-{1..6..2}-gap-{12..14..1}-suffix", 118 [ 119 "prefix-1-gap-12-suffix", 120 "prefix-1-gap-13-suffix", 121 "prefix-1-gap-14-suffix", 122 "prefix-3-gap-12-suffix", 123 "prefix-3-gap-13-suffix", 124 "prefix-3-gap-14-suffix", 125 "prefix-5-gap-12-suffix", 126 "prefix-5-gap-13-suffix", 127 "prefix-5-gap-14-suffix", 128 ], 129 None, 130 ), 131 ) 132 def test_expand_braces(self, test_case): 133 input_str, output, expected_error = test_case 134 if not expected_error: 135 self.assertEqual(output, list(utils.expand_braces(input_str))) 136 else: 137 with self.assertRaises(expected_error): 138 utils.expand_braces(input_str) 139 140 @patch("aistore.sdk.utils.parse_raw_as") 141 def test_decode_response_json(self, mock_parse): 142 response_content = "text content" 143 parsed_content = "parsed content" 144 mock_response = Mock(Response) 145 mock_response.headers = {} 146 mock_response.text = response_content 147 mock_parse.return_value = parsed_content 148 149 res = utils.decode_response(str, mock_response) 150 151 self.assertEqual(parsed_content, res) 152 mock_parse.assert_called_with(str, response_content) 153 154 def test_decode_response_msgpack(self): 155 unpacked_content = {"content key": "content value"} 156 packed_content = msgpack.encode(unpacked_content) 157 mock_response = Mock(Response) 158 mock_response.headers = {HEADER_CONTENT_TYPE: MSGPACK_CONTENT_TYPE} 159 mock_response.content = packed_content 160 161 res = utils.decode_response(dict, mock_response) 162 163 self.assertEqual(unpacked_content, res)