Updated failing Unit Tests to old PR: Add --prefer-private-ip argument to the eb ssh command #89

ebcli/controllers/ssh.py

@@ -35,6 +35,8 @@ class Meta:
             (['--setup'], dict(
                 action='store_true', help=flag_text['ssh.setup'])),
             (['--timeout'], dict(type=int, help=flag_text['ssh.timeout'])),
+            (['--prefer-private-ip'], dict(
+                action='store_true', help=flag_text['ssh.prefer_private_ip'])),
         ]
 
     def do_command(self):
@@ -47,6 +49,7 @@ def do_command(self):
         force = self.app.pargs.force
         setup = self.app.pargs.setup
         timeout = self.app.pargs.timeout
+        prefer_private_ip = self.app.pargs.prefer_private_ip
 
         if timeout and not setup:
             raise InvalidOptionsError(strings['ssh.timeout_without_setup'])
@@ -60,5 +63,6 @@ def do_command(self):
                 number=number,
                 custom_ssh=custom_ssh,
                 command=cmd,
-                timeout=timeout
+                timeout=timeout,
+                prefer_private_ip=prefer_private_ip,
         )

ebcli/operations/sshops.py

@@ -28,7 +28,8 @@
 
 def prepare_for_ssh(env_name, instance, keep_open, force, setup, number,
                     keyname=None, no_keypair_error_message=None,
-                    custom_ssh=None, command=None, timeout=None):
+                    custom_ssh=None, command=None, timeout=None,
+                    prefer_private_ip=False):
     if setup:
         setup_ssh(env_name, keyname, timeout=timeout)
         return
@@ -62,7 +63,8 @@ def prepare_for_ssh(env_name, instance, keep_open, force, setup, number,
             keep_open=keep_open,
             force_open=force,
             custom_ssh=custom_ssh,
-            command=command
+            command=command,
+            prefer_private_ip=prefer_private_ip,
         )
     except NoKeypairError:
         if not no_keypair_error_message:
@@ -85,19 +87,20 @@ def setup_ssh(env_name, keyname, timeout=None):
         commonops.update_environment(env_name, options, False, timeout=timeout or 5)
 
 
-def ssh_into_instance(instance_id, keep_open=False, force_open=False, custom_ssh=None, command=None):
+def ssh_into_instance(instance_id, keep_open=False, force_open=False, custom_ssh=None, command=None, prefer_private_ip=False):
     instance = ec2.describe_instance(instance_id)
     try:
         keypair_name = instance['KeyName']
     except KeyError:
         raise NoKeypairError()
-    try:
-        ip = instance['PublicIpAddress']
-    except KeyError:
-        if 'PrivateIpAddress' in instance:
-            ip = instance['PrivateIpAddress']
-        else:
-            raise NotFoundError(strings['ssh.noip'])
+
+    if prefer_private_ip:
+        ip = instance.get('PrivateIpAddress', instance.get('PublicIpAddress'))
+    else:
+        ip = instance.get('PublicIpAddress', instance.get('PrivateIpAddress'))
+    if ip is None:
+        raise NotFoundError(strings['ssh.noip'])
+
     security_groups = instance['SecurityGroups']
 
     user = 'ec2-user'

ebcli/resources/strings.py

@@ -874,6 +874,7 @@
     'ssh.setup': 'setup SSH for the environment',
     'ssh.timeout': "Specify the timeout period in minutes. Can only be used with the "
                    "'--setup' argument.",
+    'ssh.prefer_private_ip': "Prefer connecting to an instance's private IP address",
 
     'cleanup.resources': 'Valid values include (builder, versions, all). You can specify '
                          '"builder" to terminate the environment used to create this platform. '

tests/unit/controllers/test_ssh.py

@@ -58,7 +58,8 @@ def test_ssh(
             keep_open=False,
             number=None,
             setup=False,
-            timeout=None
+            timeout=None,
+            prefer_private_ip=False
         )
 
     @mock.patch('ebcli.controllers.ssh.SSHController.get_env_name')
@@ -83,7 +84,8 @@ def test_ssh__setup__with_timeout(
             keep_open=False,
             number=None,
             setup=True,
-            timeout=10
+            timeout=10,
+            prefer_private_ip=False
         )
 
     @mock.patch('ebcli.controllers.ssh.SSHController.get_env_name')
@@ -108,7 +110,8 @@ def test_ssh__setup__without_timeout(
             keep_open=False,
             number=None,
             setup=True,
-            timeout=None
+            timeout=None,
+            prefer_private_ip=False
         )
 
     @mock.patch('ebcli.controllers.ssh.SSHController.get_env_name')

tests/unit/operations/test_sshops.py

@@ -321,6 +321,28 @@ def test_ssh_into_instance__uses_private_address(
 
         sshops.ssh_into_instance('instance-id')
         call_mock.assert_called_once_with(['ssh', '-i', 'aws-eb-us-west-2', 'ec2-user@172.31.35.210'])
+
+    @mock.patch('ebcli.operations.sshops.ec2.describe_instance')
+    @mock.patch('ebcli.operations.sshops.ec2.describe_security_group')
+    @mock.patch('ebcli.operations.sshops.ec2.authorize_ssh')
+    @mock.patch('ebcli.operations.sshops._get_ssh_file')
+    @mock.patch('ebcli.operations.sshops.subprocess.call')
+    def test_ssh_into_instance__uses_private_address_when_private_ip_flag_is_present(
+            self,
+            call_mock,
+            _get_ssh_file_mock,
+            authorize_ssh_mock,
+            describe_security_group_mock,
+            describe_instance_mock
+    ):
+        describe_instance_response = deepcopy(mock_responses.DESCRIBE_INSTANCES_RESPONSE['Reservations'][0]['Instances'][0])
+        describe_instance_mock.return_value = describe_instance_response
+        describe_security_group_mock.return_value = mock_responses.DESCRIBE_SECURITY_GROUPS_RESPONSE['SecurityGroups'][0]
+        _get_ssh_file_mock.return_value = 'aws-eb-us-west-2'
+        call_mock.return_value = 0
+
+        sshops.ssh_into_instance('instance-id', prefer_private_ip=True)
+        call_mock.assert_called_once_with(['ssh', '-i', 'aws-eb-us-west-2', 'ec2-user@172.31.35.210'])
 
     @mock.patch('ebcli.operations.sshops.ec2.describe_instance')
     @mock.patch('ebcli.operations.sshops.ec2.describe_security_group')
@@ -520,7 +542,8 @@ def test_prepare_for_ssh__choose_instance_to_ssh_into(
             command=None,
             custom_ssh=None,
             force_open=False,
-            keep_open=False
+            keep_open=False,
+            prefer_private_ip=False
         )
 
     @mock.patch('ebcli.operations.sshops.commonops.get_instance_ids')
@@ -553,7 +576,8 @@ def test_prepare_for_ssh__number_of_instance_specified(
             command=None,
             custom_ssh=None,
             force_open=False,
-            keep_open=False
+            keep_open=False,
+            prefer_private_ip=False
         )
 
     @mock.patch('ebcli.operations.sshops.commonops.get_instance_ids')
@@ -589,7 +613,8 @@ def test_prepare_for_ssh__ssh_into_instance_fails(
             command=None,
             custom_ssh=None,
             force_open=False,
-            keep_open=False
+            keep_open=False,
+            prefer_private_ip=False
         )
         log_error_mock.assert_called_once_with(
             'This environment is not set up for SSH. Use "eb ssh --setup" to set up SSH for the environment.'