CoCalc -- test_sessionmanager.py

GitHub Repository: aws/aws-cli
Path: blob/develop/tests/unit/customizations/test_sessionmanager.py
¹⁵⁶⁷ views
1
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License"). You
4
# may not use this file except in compliance with the License. A copy of
5
# the License is located at
6
#
7
#     http://aws.amazon.com/apache2.0/
8
#
9
# or in the "license" file accompanying this file. This file is
10
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11
# ANY KIND, either express or implied. See the License for the specific
12
# language governing permissions and limitations under the License.
13
import botocore.session
14
import errno
15
import json
16
import pytest
17
import subprocess
18

19
from awscli.customizations import sessionmanager
20
from awscli.testutils import mock, unittest
21

22

23
class TestSessionManager(unittest.TestCase):
24

25
    def setUp(self):
26
        self.session = mock.Mock(botocore.session.Session)
27
        self.client = mock.Mock()
28
        self.region = 'us-west-2'
29
        self.profile = 'testProfile'
30
        self.endpoint_url = 'testUrl'
31
        self.client.meta.region_name = self.region
32
        self.client.meta.endpoint_url = self.endpoint_url
33
        self.session.create_client.return_value = self.client
34
        self.session.profile = self.profile
35
        self.caller = sessionmanager.StartSessionCaller(self.session)
36

37
        self.parsed_globals = mock.Mock()
38
        self.parsed_globals.profile = 'user_profile'
39

40
    def test_start_session_when_non_custom_start_session_fails(self):
41
        self.client.start_session.side_effect = Exception('some exception')
42
        params = {}
43
        with self.assertRaisesRegex(Exception, 'some exception'):
44
            self.caller.invoke('ssm', 'StartSession', params, mock.Mock())
45

46
    @mock.patch("awscli.customizations.sessionmanager.check_output")
47
    @mock.patch('awscli.customizations.sessionmanager.check_call')
48
    def test_start_session_success_scenario(
49
            self, mock_check_call, mock_check_output
50
    ):
51
        mock_check_output.return_value = "1.2.0.0\n"
52
        mock_check_call.return_value = 0
53

54
        start_session_params = {
55
            "Target": "i-123456789"
56
        }
57

58
        start_session_response = {
59
            "SessionId": "session-id",
60
            "TokenValue": "token-value",
61
            "StreamUrl": "stream-url"
62
        }
63

64
        self.client.start_session.return_value = start_session_response
65

66
        rc = self.caller.invoke('ssm', 'StartSession',
67
                                start_session_params, self.parsed_globals)
68
        self.assertEqual(rc, 0)
69
        self.client.start_session.assert_called_with(**start_session_params)
70

71
        mock_check_output_list = mock_check_output.call_args[0]
72
        self.assertEqual(
73
            mock_check_output_list[0], ["session-manager-plugin", "--version"]
74
        )
75

76
        mock_check_call_list = mock_check_call.call_args[0][0]
77
        mock_check_call_list[1] = json.loads(mock_check_call_list[1])
78
        self.assertEqual(
79
            mock_check_call_list,
80
            ['session-manager-plugin',
81
             start_session_response,
82
             self.region,
83
             'StartSession',
84
             self.parsed_globals.profile,
85
             json.dumps(start_session_params),
86
             self.endpoint_url]
87
        )
88

89
    @mock.patch("awscli.customizations.sessionmanager.check_output")
90
    @mock.patch('awscli.customizations.sessionmanager.check_call')
91
    def test_start_session_when_check_call_fails(
92
        self, mock_check_call, mock_check_output
93
    ):
94
        mock_check_output.return_value = "1.2.0.0\n"
95
        mock_check_call.side_effect = OSError(errno.ENOENT, 'some error')
96

97
        start_session_params = {
98
            "Target": "i-123456789"
99
        }
100

101
        start_session_response = {
102
            "SessionId": "session-id",
103
            "TokenValue": "token-value",
104
            "StreamUrl": "stream-url"
105
        }
106

107
        terminate_session_params = {
108
            "SessionId": "session-id"
109
        }
110

111
        self.client.start_session.return_value = start_session_response
112

113
        with self.assertRaises(ValueError):
114
            self.caller.invoke('ssm', 'StartSession',
115
                               start_session_params, self.parsed_globals)
116

117
            self.client.start_session.assert_called_with(
118
                **start_session_params)
119
            self.client.terminate_session.assert_called_with(
120
                **terminate_session_params)
121

122
            mock_check_call_list = mock_check_call.call_args[0][0]
123
            mock_check_call_list[1] = json.loads(mock_check_call_list[1])
124
            self.assertEqual(
125
                mock_check_call_list,
126
                ['session-manager-plugin',
127
                 start_session_response,
128
                 self.region,
129
                 'StartSession',
130
                 self.parsed_globals.profile,
131
                 json.dumps(start_session_params),
132
                 self.endpoint_url]
133
            )
134

135
    @mock.patch('awscli.customizations.sessionmanager.check_call')
136
    @mock.patch("awscli.customizations.sessionmanager.check_output")
137
    def test_start_session_when_no_profile_is_passed(
138
        self, mock_check_output, mock_check_call
139
    ):
140
        mock_check_output.return_value = "1.2.500.0\n"
141
        self.session.profile = "session_profile"
142
        self.parsed_globals.profile = None
143
        mock_check_call.return_value = 0
144

145
        start_session_params = {
146
            "Target": "i-123456789"
147
        }
148

149
        start_session_response = {
150
            "SessionId": "session-id",
151
            "TokenValue": "token-value",
152
            "StreamUrl": "stream-url"
153
        }
154
        ssm_env_name = "AWS_SSM_START_SESSION_RESPONSE"
155

156
        self.client.start_session.return_value = start_session_response
157

158
        rc = self.caller.invoke('ssm', 'StartSession',
159
                                start_session_params, self.parsed_globals)
160
        self.assertEqual(rc, 0)
161
        self.client.start_session.assert_called_with(**start_session_params)
162
        mock_check_call_list = mock_check_call.call_args[0][0]
163
        self.assertEqual(
164
            mock_check_call_list,
165
            [
166
                "session-manager-plugin",
167
                ssm_env_name,
168
                self.region,
169
                "StartSession",
170
                "",
171
                json.dumps(start_session_params),
172
                self.endpoint_url,
173
            ],
174
        )
175

176
    @mock.patch("awscli.customizations.sessionmanager.check_call")
177
    @mock.patch("awscli.customizations.sessionmanager.check_output")
178
    def test_start_session_with_env_variable_success_scenario(
179
        self, mock_check_output, mock_check_call
180
    ):
181
        mock_check_output.return_value = "1.2.500.0\n"
182
        mock_check_call.return_value = 0
183

184
        start_session_params = {
185
            "Target": "i-123456789"
186
        }
187

188
        start_session_response = {
189
            "SessionId": "session-id",
190
            "TokenValue": "token-value",
191
            "StreamUrl": "stream-url",
192
        }
193
        ssm_env_name = "AWS_SSM_START_SESSION_RESPONSE"
194

195
        self.client.start_session.return_value = start_session_response
196
        rc = self.caller.invoke(
197
            "ssm", "StartSession", start_session_params, self.parsed_globals
198
        )
199
        self.assertEqual(rc, 0)
200
        self.client.start_session.assert_called_with(**start_session_params)
201

202
        mock_check_output_list = mock_check_output.call_args[0]
203
        self.assertEqual(
204
            mock_check_output_list[0], ["session-manager-plugin", "--version"]
205
        )
206

207
        mock_check_call_list = mock_check_call.call_args[0][0]
208
        self.assertEqual(
209
            mock_check_call_list,
210
            [
211
                "session-manager-plugin",
212
                ssm_env_name,
213
                self.region,
214
                "StartSession",
215
                self.parsed_globals.profile,
216
                json.dumps(start_session_params),
217
                self.endpoint_url,
218
            ],
219
        )
220
        env_variable = mock_check_call.call_args[1]
221
        self.assertEqual(
222
            env_variable["env"][ssm_env_name],
223
            json.dumps(start_session_response)
224
        )
225

226
    @mock.patch("awscli.customizations.sessionmanager.check_call")
227
    @mock.patch("awscli.customizations.sessionmanager.check_output")
228
    def test_start_session_when_check_output_fails(
229
        self, mock_check_output, mock_check_call
230
    ):
231
        mock_check_output.side_effect = subprocess.CalledProcessError(
232
            returncode=1, cmd="session-manager-plugin", output="some error"
233
        )
234

235
        start_session_params = {
236
            "Target": "i-123456789"
237
        }
238
        start_session_response = {
239
            "SessionId": "session-id",
240
            "TokenValue": "token-value",
241
            "StreamUrl": "stream-url",
242
        }
243

244
        self.client.start_session.return_value = start_session_response
245
        with self.assertRaises(subprocess.CalledProcessError):
246
            self.caller.invoke(
247
                "ssm",
248
                "StartSession",
249
                start_session_params,
250
                self.parsed_globals
251
            )
252

253
        self.client.start_session.assert_called_with(**start_session_params)
254
        self.client.terminate_session.assert_not_called()
255
        mock_check_output.assert_called_with(
256
            ["session-manager-plugin", "--version"], text=True
257
        )
258
        mock_check_call.assert_not_called()
259

260
    @mock.patch("awscli.customizations.sessionmanager.check_call")
261
    @mock.patch("awscli.customizations.sessionmanager.check_output")
262
    def test_start_session_when_response_not_json(
263
        self, mock_check_output, mock_check_call
264
    ):
265
        mock_check_output.return_value = "1.2.500.0\n"
266
        start_session_params = {
267
            "Target": "i-123456789"
268
        }
269
        start_session_response = {
270
            "SessionId": "session-id",
271
            "TokenValue": "token-value",
272
            "StreamUrl": "stream-url",
273
            "para2": {"Not a json format"},
274
        }
275
        expected_env_value = {
276
            "SessionId": "session-id",
277
            "TokenValue": "token-value",
278
            "StreamUrl": "stream-url",
279
        }
280

281
        ssm_env_name = "AWS_SSM_START_SESSION_RESPONSE"
282

283
        self.client.start_session.return_value = start_session_response
284
        rc = self.caller.invoke(
285
            "ssm", "StartSession", start_session_params, self.parsed_globals
286
        )
287
        self.assertEqual(rc, 0)
288
        self.client.start_session.assert_called_with(**start_session_params)
289

290
        mock_check_output_list = mock_check_output.call_args[0]
291
        self.assertEqual(
292
            mock_check_output_list[0], ["session-manager-plugin", "--version"]
293
        )
294

295
        mock_check_call_list = mock_check_call.call_args[0][0]
296
        self.assertEqual(
297
            mock_check_call_list,
298
            [
299
                "session-manager-plugin",
300
                ssm_env_name,
301
                self.region,
302
                "StartSession",
303
                self.parsed_globals.profile,
304
                json.dumps(start_session_params),
305
                self.endpoint_url,
306
            ],
307
        )
308
        env_variable = mock_check_call.call_args[1]
309
        self.assertEqual(
310
            env_variable["env"][ssm_env_name], json.dumps(expected_env_value)
311
        )
312

313
    @mock.patch("awscli.customizations.sessionmanager.check_call")
314
    @mock.patch("awscli.customizations.sessionmanager.check_output")
315
    def test_start_session_when_invalid_plugin_version(
316
        self, mock_check_output, mock_check_call
317
    ):
318
        mock_check_output.return_value = "InvalidVersion"
319

320
        start_session_params = {
321
            "Target": "i-123456789"
322
        }
323
        start_session_response = {
324
            "SessionId": "session-id",
325
            "TokenValue": "token-value",
326
            "StreamUrl": "stream-url",
327
        }
328

329
        self.client.start_session.return_value = start_session_response
330
        self.caller.invoke(
331
            "ssm", "StartSession", start_session_params, self.parsed_globals
332
        )
333
        self.client.start_session.assert_called_with(**start_session_params)
334
        self.client.terminate_session.assert_not_called()
335
        mock_check_output.assert_called_with(
336
            ["session-manager-plugin", "--version"], text=True
337
        )
338

339
        mock_check_call_list = mock_check_call.call_args[0][0]
340
        self.assertEqual(
341
            mock_check_call_list,
342
            [
343
                "session-manager-plugin",
344
                json.dumps(start_session_response),
345
                self.region,
346
                "StartSession",
347
                self.parsed_globals.profile,
348
                json.dumps(start_session_params),
349
                self.endpoint_url,
350
            ],
351
        )
352

353

354
class TestVersionRequirement:
355
    version_requirement = \
356
        sessionmanager.VersionRequirement(min_version="1.2.497.0")
357

358
    @pytest.mark.parametrize(
359
        "version, expected_result",
360
        [
361
            ("2.0.0.0", True),
362
            ("2.1", True),
363
            ("2", True),
364
            ("1.3.1.1", True),
365
            ("\r\n1. 3.1.1", True),
366
            ("1.3.1.0", True),
367
            ("1.3", True),
368
            ("1.2.498.1", True),
369
            ("1.2.498", True),
370
            ("1.2.497.1", True),
371
            ("1.2.497.0", False),
372
            ("1.2.497", False),
373
            ("1.2.1.1", False),
374
            ("1.2.1", False),
375
            ("1.2", False),
376
            ("1.1.1.0", False),
377
            ("1.1.1", False),
378
            ("1.0.497.0", False),
379
            ("1.0. 497.0\r\n", False),
380
            ("1", False),
381
            ("0.3.497.0", False),
382
        ],
383
    )
384
    def test_meets_requirement(self, version, expected_result):
385
        assert expected_result == \
386
            self.version_requirement.meets_requirement(version)
387

388
    @pytest.mark.parametrize(
389
        "version, expected_result",
390
        [
391
            ("\r\n1.3.1.1", "1.3.1.1"),
392
            ("\r1.3 .1.1", "1.3.1.1"),
393
            ("1 .3.1.1", "1.3.1.1"),
394
            (" 1.3.1.1", "1.3.1.1"),
395
            ("1.3.1.1 ", "1.3.1.1"),
396
            (" 1.3.1.1 ", "1.3.1.1"),
397
            ("\n1.3.1.1 ", "1.3.1.1"),
398
            ("1.3.1.1\n", "1.3.1.1"),
399
            ("1.3\r\n.1.1", "1.3.1.1"),
400
            (" 1\r. 3", "1.3"),
401
            (" 1. 3. ", "1.3."),
402
            ("1.1.1\r\n", "1.1.1"),
403
            ("1\r", "1"),
404
        ],
405
    )
406
    def test_sanitize_plugin_version(self, version, expected_result):
407
        assert expected_result == \
408
            self.version_requirement._sanitize_plugin_version(version)
409

410
    @pytest.mark.parametrize(
411
        "version, expected_result",
412
        [
413
            ("999.99999.99.9", True),
414
            ("2", True),
415
            ("1.1.1.1", True),
416
            ("1.1.1", True),
417
            ("1.1", True),
418
            ("1.1.1.1.1", False),
419
            ("1.1.1.1.0", False),
420
            ("1.1.1.a", False),
421
            ("1.a.1.1", False),
422
            ("1-1.1.1", False),
423
            ("1.1.", False),
424
            ("invalid_version", False),
425
        ],
426
    )
427
    def test_is_valid_version(self, version, expected_result):
428
        assert expected_result == \
429
            self.version_requirement._is_valid_version(version)
430

431
Product

Resources

Company
