1 import numpy as np

2 import pytest

3

4 from muutils.json_serialize import JsonSerializer

5 from muutils.json_serialize.array import (

6 ArrayMode,

7 arr_metadata,

8 array_n_elements,

9 load_array,

10 serialize_array,

11)

12

13# pylint: disable=missing-class-docstring

14

15

16 class TestYourModule:

17 def setup_method(self):

18 self.array_1d = np.array([1, 2, 3])

19 self.array_2d = np.array([[1, 2], [3, 4]])

20 self.array_3d = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=np.int64)

21 self.array_zero_dim = np.array(42)

22 self.jser = JsonSerializer(array_mode="list")

23

24 def test_array_n_elements(self):

25 assert array_n_elements(self.array_1d) == 3

26 assert array_n_elements(self.array_2d) == 4

27 assert array_n_elements(self.array_3d) == 8

28 assert array_n_elements(self.array_zero_dim) == 1

29

30 def test_arr_metadata(self):

31 metadata = arr_metadata(self.array_3d)

32 assert metadata["shape"] == [2, 2, 2]

33 assert metadata["dtype"] == "int64"

34 assert metadata["n_elements"] == 8

35

36 @pytest.mark.parametrize(

37 "array_mode,expected_type",

38 [

39 ("list", list),

40 ("array_list_meta", dict),

41 ("array_hex_meta", dict),

42 ("array_b64_meta", dict),

43 ],

44 )

45 def test_serialize_array(self, array_mode: ArrayMode, expected_type: type):

46 result = serialize_array(

47 self.jser, self.array_2d, "test_path", array_mode=array_mode

48 )

49 assert isinstance(result, expected_type)

50

51 def test_load_array(self):

52 serialized_array = serialize_array(

53 self.jser, self.array_3d, "test_path", array_mode="array_list_meta"

54 )

55 loaded_array = load_array(serialized_array, array_mode="array_list_meta")

56 assert np.array_equal(loaded_array, self.array_3d)

57

58 def test_serialize_load_integration(self):

59 for array_mode in [

60 "list",

61 "array_list_meta",

62 "array_hex_meta",

63 "array_b64_meta",

64 ]:

65 for array in [self.array_1d, self.array_2d, self.array_3d]:

66 serialized_array = serialize_array(

67 self.jser, array, "test_path", array_mode=array_mode

68 )

69 loaded_array = load_array(serialized_array, array_mode=array_mode)

70 assert np.array_equal(loaded_array, array)

71

72 def test_serialize_load_zero_dim(self):

73 for array_mode in [

74 "list",

75 "array_list_meta",

76 "array_hex_meta",

77 "array_b64_meta",

78 ]:

79 serialized_array = serialize_array(

80 self.jser, self.array_zero_dim, "test_path", array_mode=array_mode

81 )

82 loaded_array = load_array(serialized_array)

83 assert np.array_equal(loaded_array, self.array_zero_dim)