1from unittest.mock import MagicMock 

2import pytest 

3import copy 

4 

5from pdu import PDUState 

6from drivers.apc import ApcMasterswitchPDU 

7from drivers.cyberpower import PDU41004 

8from drivers.dummy import DummyPDU 

9from drivers.snmp import BaseSnmpPDU, SnmpPDU, snmp_get, snmp_set, snmp_walk 

10 

11 

12@pytest.fixture(autouse=True) 

13def reset_easysnmp_mock(monkeypatch): 

14 import drivers.snmp 

15 global snmp_get, snmp_set, snmp_walk 

16 m1, m2, m3 = MagicMock(), MagicMock(), MagicMock() 

17 # REVIEW: I wonder if there's a clever way of covering the 

18 # difference in import locations between here and snmp.py 

19 monkeypatch.setattr(drivers.snmp, "snmp_walk", m1) 

20 monkeypatch.setattr(drivers.snmp, "snmp_get", m2) 

21 monkeypatch.setattr(drivers.snmp, "snmp_set", m3) 

22 snmp_walk = m1 

23 snmp_get = m2 

24 snmp_set = m3 

25 

26 

27def test_driver_BaseSnmpPDU_listing_ports(): 

28 pdu = BaseSnmpPDU("MyPDU", "127.0.0.1", "label_base") 

29 snmp_walk.return_value = [MagicMock(value="P1"), MagicMock(value="P2")] 

30 snmp_walk.assert_not_called() 

31 ports = pdu.ports 

32 snmp_walk.assert_called_with(pdu.oid_outlets_label_base, 

33 hostname=pdu.hostname, community=pdu.community, 

34 version=1) 

35 

36 # Check that the labels are stored, and the port IDs are 1-indexed 

37 for i in range(0, 2): 

38 assert ports[i].port_id == i+1 

39 assert ports[i].label == f"P{i+1}" 

40 

41 snmp_walk.side_effect = SystemError("An error") 

42 with pytest.raises(ValueError): 

43 pdu.ports 

44 

45 

46def test_driver_BaseSnmpPDU_port_label_mapping(): 

47 pdu = BaseSnmpPDU("MyPDU", "127.0.0.1", "label_base") 

48 pdu.port_oid = MagicMock(return_value="oid_port_1") 

49 snmp_walk.return_value = [ 

50 MagicMock(value="P1"), 

51 MagicMock(value="P2") 

52 ] 

53 snmp_set.return_value = True 

54 assert pdu.set_port_state("P1", PDUState.REBOOT) is True 

55 pdu.port_oid.assert_called_with(1) 

56 snmp_set.assert_called_with(pdu.port_oid.return_value, 

57 pdu.state_to_raw_value(PDUState.REBOOT), 'i', 

58 hostname=pdu.hostname, community=pdu.community, 

59 version=1) 

60 assert pdu.set_port_state("P2", PDUState.REBOOT) is True 

61 pdu.port_oid.assert_called_with(2) 

62 snmp_set.assert_called_with(pdu.port_oid.return_value, 

63 pdu.state_to_raw_value(PDUState.REBOOT), 'i', 

64 hostname=pdu.hostname, community=pdu.community, 

65 version=1) 

66 with pytest.raises(ValueError): 

67 pdu.set_port_state("flubberbubber", PDUState.OFF) 

68 

69 

70def test_driver_BaseSnmpPDU_get_port(): 

71 pdu = BaseSnmpPDU("MyPDU", "127.0.0.1", "label_base") 

72 

73 with pytest.raises(ValueError): 

74 pdu.port_oid(2) 

75 pdu.port_oid = MagicMock(return_value="oid_port_2") 

76 

77 snmp_get.return_value = pdu.raw_value_to_state(PDUState.REBOOT) 

78 snmp_get.assert_not_called() 

79 assert pdu.get_port_state(2) 

80 pdu.port_oid.assert_called_with(2) 

81 snmp_get.assert_called_with(pdu.port_oid.return_value, 

82 hostname=pdu.hostname, community=pdu.community, 

83 version=1) 

84 

85 snmp_get.side_effect = SystemError("An error") 

86 with pytest.raises(ValueError): 

87 pdu.get_port_state(2) 

88 

89 

90def test_driver_BaseSnmpPDU_set_port(): 

91 pdu = BaseSnmpPDU("MyPDU", "127.0.0.1", "label_base") 

92 

93 pdu.port_oid = MagicMock(return_value="oid_port_2") 

94 snmp_set.return_value = True 

95 snmp_set.assert_not_called() 

96 assert pdu.set_port_state(2, PDUState.REBOOT) is True 

97 pdu.port_oid.assert_called_with(2) 

98 snmp_set.assert_called_with(pdu.port_oid.return_value, 

99 pdu.state_to_raw_value(PDUState.REBOOT), 'i', 

100 hostname=pdu.hostname, community=pdu.community, 

101 version=1) 

102 

103 snmp_set.side_effect = SystemError("An error") 

104 with pytest.raises(ValueError): 

105 pdu.set_port_state(2, PDUState.REBOOT) 

106 

107 

108def test_driver_BaseSnmpPDU_action_translation(): 

109 pdu = BaseSnmpPDU("MyPDU", "127.0.0.1", "label_base") 

110 

111 # Check the state -> SNMP value translation 

112 for action in PDUState.valid_actions(): 

113 assert pdu.state_to_raw_value(action) == pdu.action_to_snmp_value[action.name] 

114 

115 with pytest.raises(ValueError): 

116 pdu.state_to_raw_value(PDUState.UNKNOWN) 

117 

118 # Check the SNMP value -> state translation 

119 for state in PDUState.valid_actions(): 

120 raw = pdu.state_to_raw_value(state) 

121 assert pdu.raw_value_to_state(raw) == state 

122 

123 with pytest.raises(AttributeError): 

124 pdu.state_to_raw_value(42) 

125 

126 

127def test_driver_ApcMasterswitchPDU_check_OIDs(): 

128 pdu = ApcMasterswitchPDU("MyPDU", config={"hostname": "127.0.0.1"}) 

129 

130 assert pdu.oid_outlets_label_base == "SNMPv2-SMI::enterprises.318.1.1.4.4.2.1.4" 

131 assert pdu.port_oid(10) == "SNMPv2-SMI::enterprises.318.1.1.4.4.2.1.3.10" 

132 

133 

134def test_driver_ApcMasterswitchPDU_invalid_config(): 

135 with pytest.raises(ValueError): 

136 ApcMasterswitchPDU("MyPDU", config={}) 

137 

138 

139def test_driver_PDU41004_check_OIDs(): 

140 pdu = PDU41004("MyPDU", config={"hostname": "127.0.0.1"}) 

141 

142 assert pdu.oid_outlets_label_base == "SNMPv2-SMI::enterprises.3808.1.1.3.3.3.1.1.2" 

143 assert pdu.port_oid(10) == "SNMPv2-SMI::enterprises.3808.1.1.3.3.3.1.1.4.10" 

144 

145 

146def test_driver_PDU41004_invalid_config(): 

147 with pytest.raises(ValueError): 

148 PDU41004("MyPDU", config={}) 

149 

150 

151def test_driver_DummyPDU(): 

152 ports = ['P1', 'P2', 'P3'] 

153 pdu = DummyPDU("MyPDU", {"ports": ports}) 

154 

155 assert [p.label for p in pdu.ports] == ports 

156 assert pdu.get_port_state(0) == PDUState.ON 

157 pdu.set_port_state(0, PDUState.OFF) 

158 assert pdu.get_port_state(0) == PDUState.OFF 

159 

160 

161def test_driver_SnmpPDU_check_OIDs_and_default_actions(): 

162 pdu = SnmpPDU("MyPDU", config={ 

163 "hostname": "127.0.0.1", 

164 "oid_outlets_label_base": "label_base", 

165 "oid_outlets_base": "outlet_base", 

166 }) 

167 

168 assert pdu.community == "private" 

169 assert pdu.oid_outlets_label_base == "label_base" 

170 assert pdu.port_oid(10) == "outlet_base.10" 

171 assert pdu.action_to_snmp_value == super(SnmpPDU, pdu).action_to_snmp_value 

172 

173 

174def test_driver_SnmpPDU_actions(): 

175 pdu = SnmpPDU("MyPDU", config={ 

176 "hostname": "127.0.0.1", 

177 "oid_outlets_label_base": "label_base", 

178 "oid_outlets_base": "outlet_base", 

179 "community": "public", 

180 "action_to_snmp_value": { 

181 "ON": 42, 

182 "OFF": 43, 

183 "REBOOT": 44 

184 } 

185 }) 

186 

187 assert pdu.community == "public" 

188 assert pdu.action_to_snmp_value == { 

189 "ON": 42, 

190 "OFF": 43, 

191 "REBOOT": 44 

192 } 

193 

194 

195def test_driver_SnmpPDU_invalid_actions(): 

196 with pytest.raises(ValueError): 

197 SnmpPDU("MyPDU", config={ 

198 "hostname": "127.0.0.1", 

199 "oid_outlets_label_base": "label_base", 

200 "oid_outlets_base": "outlet_base", 

201 "community": "public", 

202 "action_to_snmp_value": { 

203 "ON": "TOTO", 

204 "OFF": 43, 

205 "REBOOT": 44 

206 } 

207 }) 

208 

209 

210def test_driver_SnmpPDU_missing_actions(): 

211 with pytest.raises(ValueError): 

212 SnmpPDU("MyPDU", config={ 

213 "hostname": "127.0.0.1", 

214 "oid_outlets_label_base": "label_base", 

215 "oid_outlets_base": "outlet_base", 

216 "community": "public", 

217 "action_to_snmp_value": { 

218 "OFF": 43, 

219 "REBOOT": 44 

220 } 

221 }) 

222 

223 

224def test_driver_SnmpPDU_missing_parameters(): 

225 valid_config = { 

226 "hostname": "127.0.0.1", 

227 "oid_outlets_label_base": "label_base", 

228 "oid_outlets_base": "outlet_base", 

229 "community": "public", 

230 "action_to_snmp_value": { 

231 "ON": 42, 

232 "OFF": 43, 

233 "REBOOT": 44 

234 } 

235 } 

236 

237 SnmpPDU("MyPDU", config=valid_config) 

238 for required_param in ["hostname", "oid_outlets_label_base", "oid_outlets_base"]: 

239 new_config = copy.deepcopy(valid_config) 

240 del new_config[required_param] 

241 with pytest.raises(ValueError): 

242 SnmpPDU("MyPDU", config=new_config)