Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
aws
GitHub Repository: aws/aws-cli
 Path: blob/develop/tests/unit/customizations/emr/test_sshutils.py
1569 views
1
# Copyright 2014 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


14
from awscli.customizations.emr import sshutils

15
from awscli.customizations.emr import exceptions

16
from awscli.testutils import mock, unittest

17


18


19
class TestSSHUtils(unittest.TestCase):

20


21
    @mock.patch('awscli.customizations.emr.sshutils.emrutils')

22
    def test_validate_and_find_master_dns_waits(self, emrutils):

23
        emrutils.get_cluster_state.return_value = 'STARTING'

24
        session = mock.Mock()

25
        client = mock.Mock()

26
        emrutils.get_client.return_value = client

27


28
        sshutils.validate_and_find_master_dns(session, None, 'cluster-id')

29


30
        # We should have:

31
        # 1. Waiter for the cluster to be running.

32
        client.get_waiter.assert_called_with('cluster_running')

33
        client.get_waiter.return_value.wait.assert_called_with(

34
            ClusterId='cluster-id')

35


36
        # 2. Found the master public DNS

37
        self.assertTrue(emrutils.find_master_dns.called)

38


39
    @mock.patch('awscli.customizations.emr.sshutils.emrutils')

40
    def test_cluster_in_terminated_states(self, emrutils):

41
        emrutils.get_cluster_state.return_value = 'TERMINATED'

42
        with self.assertRaises(exceptions.ClusterTerminatedError):

43
            sshutils.validate_and_find_master_dns(

44
                mock.Mock(), None, 'cluster-id')

45


46
    @mock.patch('awscli.customizations.emr.sshutils.emrutils')

47
    def test_ssh_scp_key_file_format(self, emrutils):

48
        def which_side_effect(program):

49
            if program == 'ssh' or program == 'scp':

50
                return '/some/path'

51
        emrutils.which.side_effect = which_side_effect

52


53
        key_file1 = 'key.abc'

54
        sshutils.validate_ssh_with_key_file(key_file1)

55
        sshutils.validate_scp_with_key_file(key_file1)

56


57
        key_file2 = 'key'

58
        sshutils.validate_ssh_with_key_file(key_file2)

59
        sshutils.validate_scp_with_key_file(key_file2)

60


61