Coverage for NeuralTSNE/NeuralTSNE/Plotter/tests/test_plot.py: 100%

1import io 

2from unittest.mock import MagicMock, call, mock_open, patch 

3 

4import numpy as np 

5import pytest 

6 

7import NeuralTSNE.Plotter as plotter 

8 

9 

10@pytest.mark.parametrize("are_neural_labels", [True, False]) 

11@pytest.mark.parametrize("file_step", [None, 2]) 

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

13@patch("matplotlib.pyplot.subplots") 

14@patch("matplotlib.pyplot.xlabel") 

15@patch("matplotlib.pyplot.ylabel") 

16@patch("matplotlib.pyplot.scatter") 

17@patch("matplotlib.pyplot.savefig") 

18@patch("matplotlib.pyplot.show") 

19def test_plot( 

20 mock_plt_show: MagicMock, 

21 mock_plt_savefig: MagicMock, 

22 mock_plt_scatter: MagicMock, 

23 mock_plt_ylabel: MagicMock, 

24 mock_plt_xlabel: MagicMock, 

25 mock_plt_subplots: MagicMock, 

26 step: int, 

27 file_step: int | None, 

28 are_neural_labels: bool, 

29): 

30 img_file = "test.png" 

31 mock_plt_subplots.return_value = (None, None) 

32 plotter.plot( 

33 np.array([[1, 2], [3, 4], [5, 6], [7, 8]]), 

34 np.array([1, 2, 3, 4]), 

35 step, 

36 2, 

37 1, 

38 are_neural_labels, 

39 img_file, 

40 {"file_step": file_step} if file_step else None, 

41 ) 

42 

43 mock_plt_scatter.assert_called_once() 

44 mock_plt_savefig.assert_called_once_with(img_file) 

45 mock_plt_show.assert_called_once() 

46 mock_plt_xlabel.assert_called_once() 

47 mock_plt_ylabel.assert_called_once() 

48 mock_plt_subplots.assert_called_once() 

49 

50 

51@pytest.mark.parametrize("input_file", ["test_data.txt"]) 

52@pytest.mark.parametrize("labels_file", ["test_labels.txt", None]) 

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

54@pytest.mark.parametrize("marker_size", [1, 3]) 

55@pytest.mark.parametrize("alpha", [1, 0.5]) 

56@pytest.mark.parametrize("columns", [None, 2]) 

57@pytest.mark.parametrize("are_neural_labels", [True, False]) 

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

59@patch("numpy.loadtxt") 

60@patch("NeuralTSNE.Plotter.plot.plot") 

61def test_plot_from_file( 

62 mock_plot: MagicMock, 

63 mock_loadtxt: MagicMock, 

64 mock_files: MagicMock, 

65 are_neural_labels: bool, 

66 columns: int | None, 

67 alpha: float, 

68 marker_size: float, 

69 step: int, 

70 labels_file: str | None, 

71 input_file: str, 

72): 

73 file_content = "20\n1 2 3\n4 5 6\n7 8 9\n10 11 12" 

74 labels_content = "1 4\n2 5\n3 6" 

75 

76 handlers = ( 

77 [ 

78 io.StringIO(file_content), 

79 io.StringIO(labels_content), 

80 ] 

81 if labels_file 

82 else [io.StringIO(file_content)] 

83 ) 

84 

85 mock_files.side_effect = handlers 

86 

87 data_list = list(line.split() for line in file_content.splitlines()) 

88 data = np.array(data_list[1:], dtype="float32") 

89 

90 labels_list = list(line.split() for line in labels_content.splitlines()) 

91 labels = np.array(labels_list, dtype="int32") if labels_file else None 

92 

93 mock_loadtxt.side_effect = [data, labels] 

94 

95 file_step = int(data_list[0][0]) 

96 if labels_file: 

97 data = data[: len(labels)] 

98 

99 plotter.plot_from_file( 

100 input_file, 

101 labels_file, 

102 columns, 

103 step, 

104 marker_size, 

105 alpha, 

106 are_neural_labels, 

107 ) 

108 

109 if labels_file: 

110 mock_loadtxt.assert_has_calls( 

111 [ 

112 call(handlers[0]), 

113 call(handlers[1], usecols=columns, dtype="int"), 

114 ], 

115 any_order=True, 

116 ) 

117 mock_files.assert_has_calls( 

118 [call(input_file, "r"), call(labels_file, "r")], any_order=True 

119 ) 

120 else: 

121 mock_loadtxt.assert_called_once_with(handlers[0]) 

122 mock_files.assert_called_once_with(input_file, "r") 

123 

124 mock_plot.assert_called_once() 

125 

126 mock_plot_args = mock_plot.call_args[0] 

127 np.testing.assert_array_equal(mock_plot_args[0], data) 

128 np.testing.assert_array_equal(mock_plot_args[1], labels) 

129 assert mock_plot_args[2] == step 

130 assert mock_plot_args[3] == marker_size 

131 assert mock_plot_args[4] == alpha 

132 assert mock_plot_args[5] == are_neural_labels 

133 assert mock_plot_args[6] == input_file.rsplit(".", 1)[0] + ".png" 

134 assert mock_plot_args[7] == {"file_step": file_step}