Coverage for NeuralTSNE/NeuralTSNE/Utils/Loaders/FileLoaders/tests/test_file_loaders.py: 100%

1import io 

2from typing import List 

3from unittest.mock import MagicMock, mock_open, patch 

4 

5import numpy as np 

6import pytest 

7import torch 

8 

9from NeuralTSNE.Utils.Loaders.FileLoaders import ( 

10 load_npy_file, 

11 load_text_file, 

12 load_torch_dataset, 

13) 

14 

15 

16@pytest.mark.parametrize( 

17 "file_content, header", 

18 [ 

19 ("1\t3\t7\t8\n4\t1\t3\t2\n7\t1\t9\t5\n3\t1\t12\t9\n", False), 

20 ( 

21 "col1\tcol2\tcol3\tcol4\n1\t3\t7\t8\n4\t1\t3\t2\n7\t1\t9\t5\n3\t1\t12\t9\n", 

22 True, 

23 ), 

24 ("1 3 7 8\n4 1 3 2\n7 1 9 5\n3 1 12 9\n", False), 

25 ("col1 col2 col3 col4\n1 3 7 8\n4 1 3 2\n7 1 9 5\n3 1 12 9\n", True), 

26 ], 

27) 

28@pytest.mark.parametrize("variance_threshold", [0.1, 0.5, 0.9]) 

29@pytest.mark.parametrize("exclude_cols", [[3], [1, 2], [3, 1], None]) 

30@pytest.mark.parametrize("step", [1, 2, 3]) 

31@patch("NeuralTSNE.Utils.Loaders.FileLoaders.file_loaders.prepare_data") 

32@patch("builtins.open", new_callable=mock_open) 

33def test_load_text_file( 

34 mock_open: MagicMock, 

35 mock_prepare_data: MagicMock, 

36 file_content: str, 

37 step: int, 

38 header: bool, 

39 exclude_cols: List[int] | None, 

40 variance_threshold: float, 

41): 

42 mock_content = io.StringIO(file_content) 

43 mock_open.return_value = mock_content 

44 start_index = 1 if header else 0 

45 matrix = np.array( 

46 [ 

47 [float(value) for value in line.split()] 

48 for line in file_content.splitlines()[start_index:] 

49 ] 

50 ) 

51 

52 if exclude_cols is not None: 

53 cols = [col for col in range(len(matrix[0])) if col not in exclude_cols] 

54 matrix = matrix[:, cols] 

55 

56 matrix = matrix[::step] 

57 matrix_t = torch.tensor(matrix).T 

58 mock_prepare_data.return_value = matrix_t 

59 

60 result = load_text_file(mock_open, step, header, exclude_cols, variance_threshold) 

61 

62 mock_prepare_data.assert_called_once() 

63 prepare_data_args = mock_prepare_data.call_args[0] 

64 assert mock_content.closed 

65 assert prepare_data_args[0] == variance_threshold 

66 assert np.allclose(prepare_data_args[1], matrix) 

67 assert torch.allclose(result, matrix_t) 

68 

69 

70@pytest.mark.parametrize( 

71 "file_content", 

72 ["1 2 3 4\n5 6 7 8\n9 10 11 12\n", "1\t2\t3\t4\n5\t6\t7\t8\n9\t10\t11\t12\n"], 

73) 

74@pytest.mark.parametrize("step", [1, 2, 3]) 

75@pytest.mark.parametrize("exclude_cols", [[3], [1, 2], [3, 1], None]) 

76@pytest.mark.parametrize("variance_threshold", [0.1, 0.5, 0.9]) 

77@patch("NeuralTSNE.Utils.Loaders.FileLoaders.file_loaders.prepare_data") 

78def test_load_npy_file( 

79 mock_prepare_data: MagicMock, 

80 file_content: str, 

81 step: int, 

82 exclude_cols: List[int] | None, 

83 variance_threshold: float, 

84): 

85 matrix = np.loadtxt(io.StringIO(file_content)) 

86 

87 matrix_file = io.BytesIO() 

88 np.save(matrix_file, matrix) 

89 matrix_file.seek(0) 

90 

91 if exclude_cols is not None: 

92 cols = [col for col in range(len(matrix[0])) if col not in exclude_cols] 

93 matrix = matrix[:, cols] 

94 

95 matrix = matrix[::step] 

96 matrix_t = torch.tensor(matrix).T 

97 mock_prepare_data.return_value = matrix_t 

98 

99 result = load_npy_file(matrix_file, step, exclude_cols, variance_threshold) 

100 

101 mock_prepare_data.assert_called_once() 

102 prepare_data_args = mock_prepare_data.call_args[0] 

103 assert prepare_data_args[0] == variance_threshold 

104 assert np.allclose(prepare_data_args[1], matrix) 

105 assert torch.allclose(result, matrix_t) 

106 

107 

108@pytest.mark.parametrize("data_name, step, output_path", [("test", 2, "output")]) 

109@patch("torch.stack") 

110@patch("NeuralTSNE.Utils.Loaders.FileLoaders.file_loaders.save_torch_labels") 

111@patch("NeuralTSNE.Utils.Loaders.FileLoaders.file_loaders.save_means_and_vars") 

112@patch("NeuralTSNE.Utils.Loaders.FileLoaders.file_loaders.get_datasets.get_dataset") 

113def test_load_torch_dataset( 

114 mock_get_dataset: MagicMock, 

115 mock_save_means_and_vars: MagicMock, 

116 mock_save_torch_labels: MagicMock, 

117 mock_stack: MagicMock, 

118 data_name: str, 

119 step: int, 

120 output_path: str, 

121): 

122 train_data = torch.tensor([1, 2, 3]) 

123 test_data = torch.tensor([1, 2, 3]) 

124 train_data = torch.utils.data.TensorDataset(train_data) 

125 test_data = torch.utils.data.TensorDataset(test_data) 

126 mock_get_dataset.return_value = train_data, test_data 

127 

128 load_torch_dataset(data_name, step, output_path) 

129 

130 mock_get_dataset.assert_called_once_with(data_name) 

131 mock_stack.assert_called_once() 

132 mock_save_torch_labels.assert_called_once_with(output_path, test_data) 

133 mock_save_means_and_vars.assert_called_once_with(mock_stack.return_value)