Coverage for tests/unit/no_torch/test_load_item_recursive.py: 99%

1 from __future__ import annotations

2

3 import typing

4 from pathlib import Path

5

6 import numpy as np

7 import pytest

8 from muutils.errormode import ErrorMode

9 from muutils.json_serialize import (

10 SerializableDataclass,

11 serializable_dataclass,

12 serializable_field,

13)

14 from muutils.json_serialize.util import _FORMAT_KEY

15

16 from zanj import ZANJ

17 from zanj.loading import LoadedZANJ, load_item_recursive

18

19 TEST_DATA_PATH: Path = Path("tests/junk_data")

20

21

22 def test_load_item_recursive_basic():

23 """Test basic functionality of load_item_recursive"""

24 # Simple JSON data

25 json_data = {

26 "name": "test",

27 "value": 42,

28 "list": [1, 2, 3],

29 "nested": {"a": 1, "b": 2},

30 }

31

32 # Load with default parameters

33 result = load_item_recursive(json_data, tuple(), None)

34

35 # Check the result

36 assert result == json_data

37 assert result["name"] == "test"

38 assert result["value"] == 42

39 assert result["list"] == [1, 2, 3]

40 assert result["nested"] == {"a": 1, "b": 2}

41

42

43 def test_load_item_recursive_numpy_array():

44 """Test loading a numpy array"""

45 # Create a JSON representation of a numpy array properly formatted

46 array_data = np.random.rand(5, 5)

47 json_data = {

48 _FORMAT_KEY: "numpy.ndarray:array_list_meta", # Use the correct format suffix

49 "dtype": str(array_data.dtype),

50 "shape": list(array_data.shape),

51 "data": array_data.tolist(),

52 }

53

54 # Load with default parameters

55 result = load_item_recursive(json_data, tuple(), None)

56

57 # Check the result

58 assert isinstance(result, np.ndarray)

59 assert result.shape == tuple(json_data["shape"])

60 assert result.dtype == np.dtype(json_data["dtype"])

61 assert np.allclose(result, array_data)

62

63

64 def test_load_item_recursive_serializable_dataclass():

65 """Test loading a SerializableDataclass"""

66

67 @serializable_dataclass

68 class TestClass(SerializableDataclass):

69 name: str

70 value: int

71 data: typing.List[int] = serializable_field(default_factory=list)

72

73 # Create an instance and serialize it

74 instance = TestClass("test", 42, [1, 2, 3])

75 serialized = instance.serialize()

76

77 # Load with default parameters

78 result = load_item_recursive(serialized, tuple(), None)

79

80 # Check the result

81 assert isinstance(result, TestClass)

82 assert result.name == "test"

83 assert result.value == 42

84 assert result.data == [1, 2, 3]

85

86

87 def test_load_item_recursive_nested_container():

88 """Test loading with nested containers"""

89 # Create a complex nested structure with properly formatted arrays

90 json_data = {

91 "name": "test",

92 "arrays": [

93 {

94 _FORMAT_KEY: "numpy.ndarray:array_list_meta", # Use correct format suffix

95 "dtype": "float64",

96 "shape": [3, 3],

97 "data": np.random.rand(3, 3).tolist(),

98 },

99 {

100 _FORMAT_KEY: "numpy.ndarray:array_list_meta", # Use correct format suffix

101 "dtype": "float64",

102 "shape": [2, 2],

103 "data": np.random.rand(2, 2).tolist(),

104 },

105 ],

106 "nested": {

107 "dict_with_array": {

108 _FORMAT_KEY: "numpy.ndarray:array_list_meta", # Use correct format suffix

109 "dtype": "float64",

110 "shape": [4, 4],

111 "data": np.random.rand(4, 4).tolist(),

112 }

113 },

114 }

115

116 # Load with default parameters

117 result = load_item_recursive(json_data, tuple(), None)

118

119 # Check the result

120 assert result["name"] == "test"

121 assert len(result["arrays"]) == 2

122 assert isinstance(result["arrays"][0], np.ndarray)

123 assert isinstance(result["arrays"][1], np.ndarray)

124 assert result["arrays"][0].shape == (3, 3)

125 assert result["arrays"][1].shape == (2, 2)

126 assert isinstance(result["nested"]["dict_with_array"], np.ndarray)

127 assert result["nested"]["dict_with_array"].shape == (4, 4)

128

129

130 def test_load_item_recursive_unknown_format():

131 """Test loading with an unknown format key"""

132 # Create JSON data with an unknown format that is not registered in the handlers

133 json_data = {

134 _FORMAT_KEY: "unknown.format.that.definitely.does.not.exist",

135 "data": [1, 2, 3],

136 }

137

138 # Load with default parameters (should return the JSON as is)

139 result = load_item_recursive(json_data, tuple(), None, allow_not_loading=True)

140

141 # Check the result

142 assert result == json_data

143

144 # TODO: this doesn't raise any errors

145 # Test with allow_not_loading=False (should raise an error)

146 # Create a ZANJ with EXCEPT error mode to ensure value errors are raised

147 z = ZANJ(error_mode=ErrorMode.EXCEPT)

148 load_item_recursive(

149 json_data, tuple(), z, error_mode=ErrorMode.EXCEPT, allow_not_loading=False

150 )

151

152

153 def test_load_item_recursive_with_external_reference():

154 """Test loading an item with an external reference"""

155 # Create a ZANJ object and save some data to create externals

156 z = ZANJ(external_array_threshold=10)

157 data = {"large_array": np.random.rand(20, 20)}

158 path = TEST_DATA_PATH / "test_load_item_recursive_external.zanj"

159 z.save(data, path)

160

161 # Load the ZANJ file

162 loaded_zanj = LoadedZANJ(path, z)

163

164 # Try loading the data

165 loaded_zanj.populate_externals()

166

167 # Check that the externals were populated

168 assert len(loaded_zanj._externals) > 0

169

170 # Verify JSON data structure

171 assert "_REF_KEY" in loaded_zanj._json_data or isinstance(

172 loaded_zanj._json_data, dict

173 )

174

175

176 def test_load_item_recursive_error_modes():

177 """Test different error modes"""

178 # Create JSON data with an unknown format

179 json_data = {

180 _FORMAT_KEY: "unknown.format.that.definitely.does.not.exist",

181 "data": [1, 2, 3],

182 }

183

184 # Test WARN mode (should not raise, just return the data)

185 result = load_item_recursive(

186 json_data, tuple(), None, error_mode=ErrorMode.WARN, allow_not_loading=True

187 )

188 assert result == json_data

189

190 # Test IGNORE mode (should not raise, just return the data)

191 result = load_item_recursive(

192 json_data, tuple(), None, error_mode=ErrorMode.IGNORE, allow_not_loading=True

193 )

194 assert result == json_data

195

196 # Create a custom class that's known to fail during loading

197 class CustomHandler:

198 def check(self, json_item, path=None, z=None):

199 return (

200 json_item.get(_FORMAT_KEY)

201 == "unknown.format.that.definitely.does.not.exist"

202 )

203

204 def load(self, json_item, path=None, z=None):

205 # This will raise a ValueError

206 raise ValueError("Forced error for testing purposes")

207

208 # Register this handler temporarily

209 import zanj.loading

210

211 original_get_item_loader = zanj.loading.get_item_loader

212

213 def mock_get_item_loader(*args, **kwargs):

214 # Always return our custom handler

215 return CustomHandler()

216

217 try:

218 # Override the get_item_loader function

219 zanj.loading.get_item_loader = mock_get_item_loader

220

221 # Test EXCEPT mode (should raise)

222 with pytest.raises(ValueError):

223 load_item_recursive(

224 json_data,

225 tuple(),

226 None,

227 error_mode=ErrorMode.EXCEPT,

228 allow_not_loading=True,

229 )

230 finally:

231 # Restore the original function

232 zanj.loading.get_item_loader = original_get_item_loader