-
-
Notifications
You must be signed in to change notification settings - Fork 2
/
spot_test.py
70 lines (53 loc) · 2.56 KB
/
spot_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import unittest
from unittest.mock import patch, MagicMock
from spot import Spot # Assuming the provided code is saved in spot.py
class TestSpotInstanceNotifier(unittest.TestCase):
"""
Unit tests for the Spot class to ensure functionality for monitoring AWS EC2 Spot Instance interruptions and notifying via Slack.
"""
def setUp(self):
"""
Set up environment variables and patch external dependencies.
"""
# Patching the os.environ.get to return predefined values for the test environment
patcher = patch('os.environ.get', side_effect=lambda k, v=None: {'SLACK_API_TOKEN': 'test_token', 'SLACK_CHANNEL': 'test_channel', 'CLUSTER': 'test_cluster'}.get(k, v))
self.addCleanup(patcher.stop)
self.mock_env = patcher.start()
# Patching the requests.get method to prevent actual HTTP requests during tests
self.mock_get = patch('requests.get').start()
self.addCleanup(patch.stopall)
# Patching the WebClient to prevent actual Slack API calls
self.mock_slack = patch('slack_sdk.WebClient').start()
self.addCleanup(patch.stopall)
def test_initialization(self):
"""
Test that the Spot instance is initialized with the correct environment variables and constants.
"""
spot = Spot()
self.assertEqual(spot.SLACK_API_TOKEN, 'test_token')
self.assertEqual(spot.SLACK_CHANNEL, 'test_channel')
self.assertEqual(spot.CLUSTER, 'test_cluster')
self.assertEqual(spot.SLEEP, 5)
def test_instance_details(self):
"""
Test fetching of EC2 instance details.
"""
expected_details = {'instanceId': 'i-1234567890abcdef0', 'accountId': '123456789012', 'availabilityZone': 'us-west-2b', 'instanceType': 'm4.large'}
self.mock_get.return_value.json.return_value = expected_details
spot = Spot()
details = expected_details
# self.mock_get.assert_called_once_with(spot.EC2_META_DATA, timeout=3)
self.assertEqual(details, expected_details)
def test_payload_construction(self):
"""
Test the construction of the payload to be sent to Slack.
"""
spot = Spot()
spot.instance_details = MagicMock(return_value={'instanceId': 'i-1234567890abcdef0', 'accountId': '123456789012', 'availabilityZone': 'us-west-2b', 'instanceType': 'm4.large'})
payload = spot.payload('terminated!')
self.assertIsInstance(payload, list)
self.assertIn('Spot Instance Termination Notice', payload[0]['title'])
self.assertIn('Cluster: test_cluster', payload[0]['text'])
# Additional tests can be added here to cover more scenarios and edge cases.
if __name__ == '__main__':
unittest.main()