Merge pull request #3 from jneprz/personal/johnli1/test_coverage

Personal/johnli1/test coverage

Author: jneprz

src/sftp/azext_sftp/_help.py

@@ -46,11 +46,11 @@
         The certificate can be used with 'az sftp connect' or with standard SFTP clients.
     examples:
         - name: Generate a certificate using an existing public key
-          text: az sftp cert --public-key-file ~/.ssh/id_rsa.pub --output-file ~/my_cert.pub
+          text: az sftp cert --public-key-file ~/.ssh/id_rsa.pub --file ~/my_cert.pub
         - name: Generate a certificate and create a new key pair in the same directory
-          text: az sftp cert --output-file ~/my_cert.pub
+          text: az sftp cert --file ~/my_cert.pub
         - name: Generate a certificate with custom SSH client folder
-          text: az sftp cert --output-file ~/my_cert.pub --ssh-client-folder "C:\\Program Files\\OpenSSH"
+          text: az sftp cert --file ~/my_cert.pub --ssh-client-folder "C:\\Program Files\\OpenSSH"
 """
 
 helps['sftp connect'] = """

src/sftp/azext_sftp/_params.py

@@ -8,11 +8,11 @@
 def load_arguments(self, _):
 
     with self.argument_context('sftp cert') as c:
-        c.argument('cert_path', options_list=['--output-file', '-o'],
+        c.argument('cert_path', options_list=['--file', '-f'],
                    help='The file path to write the SSH cert to, defaults to public key path with -aadcert.pub appended')
         c.argument('public_key_file', options_list=['--public-key-file', '-p'],
                    help='The RSA public key file path. If not provided, '
-                   'generated key pair is stored in the same directory as --output-file.')
+                   'generated key pair is stored in the same directory as --file.')
         c.argument('ssh_client_folder', options_list=['--ssh-client-folder'],
                    help='Folder path that contains ssh executables (ssh-keygen, ssh). '
                    'Default to ssh executables in your PATH or C:\\Windows\\System32\\OpenSSH on Windows.')

src/sftp/azext_sftp/custom.py

@@ -24,7 +24,7 @@ def sftp_cert(cmd, cert_path=None, public_key_file=None, ssh_client_folder=None)
     logger.debug("Starting SFTP certificate generation")
 
     if not cert_path and not public_key_file:
-        raise azclierror.RequiredArgumentMissingError("--output-file or --public-key-file must be provided.")
+        raise azclierror.RequiredArgumentMissingError("--file or --public-key-file must be provided.")
 
     if cert_path:
         cert_path = os.path.expanduser(cert_path)
@@ -52,14 +52,14 @@ def sftp_cert(cmd, cert_path=None, public_key_file=None, ssh_client_folder=None)
         logger.debug("Will generate key pair in: %s", keys_folder)
 
     try:
-        public_key_file, _, _ = _check_or_create_public_private_files(
+        public_key_file, _, delete_keys = _check_or_create_public_private_files(
             public_key_file, None, keys_folder, ssh_client_folder)
         cert_file, _ = _get_and_write_certificate(cmd, public_key_file, cert_path, ssh_client_folder)
     except Exception as e:
         logger.debug("Certificate generation failed: %s", str(e))
         raise
 
-    if keys_folder:
+    if keys_folder and delete_keys:
         logger.warning("%s contains sensitive information (id_rsa, id_rsa.pub). "
                        "Please delete once this certificate is no longer being used.", keys_folder)
 

src/sftp/azext_sftp/file_utils.py

@@ -96,15 +96,22 @@ def check_or_create_public_private_files(public_key_file, private_key_file, cred
     delete_keys = False
 
     if not public_key_file and not private_key_file:
-        delete_keys = True
         if not credentials_folder:
             credentials_folder = tempfile.mkdtemp(prefix="aadsftpcert")
         else:
             if not os.path.isdir(credentials_folder):
                 os.makedirs(credentials_folder)
+
         public_key_file = os.path.join(credentials_folder, "id_rsa.pub")
         private_key_file = os.path.join(credentials_folder, "id_rsa")
-        sftp_utils.create_ssh_keyfile(private_key_file, ssh_client_folder)
+
+        # Check if existing keys are present before generating new ones
+        if not (os.path.isfile(public_key_file) and os.path.isfile(private_key_file)):
+            # Only generate new keys if both don't exist
+            sftp_utils.create_ssh_keyfile(private_key_file, ssh_client_folder)
+            # Only set delete_keys to True when we actually create new keys
+            delete_keys = True
+        # If existing keys are found, delete_keys remains False
 
     if not public_key_file:
         if private_key_file:
@@ -155,7 +162,7 @@ def get_and_write_certificate(cmd, public_key_file, cert_file, ssh_client_folder
     telemetry.add_extension_event('sftp', {'Context.Default.AzureCLI.SFTPGetCertificateTime': time_elapsed})
 
     if not cert_file:
-        cert_file = str(public_key_file) + "-aadcert.pub"
+        cert_file = str(public_key_file.removesuffix(".pub")) + "-aadcert.pub"
 
     logger.debug("Generating certificate %s", cert_file)
     _write_cert_file(certificate, cert_file)

src/sftp/azext_sftp/sftp_utils.py

@@ -175,6 +175,9 @@ def get_ssh_cert_info(cert_file, ssh_client_folder=None):
                                          const.RECOMMENDATION_SSH_CLIENT_NOT_FOUND)
 
 
+_warned_ssh_client_folders = set()
+
+
 def get_ssh_client_path(ssh_command="ssh", ssh_client_folder=None):
     """Get the path to an SSH client executable."""
     if ssh_client_folder:
@@ -184,8 +187,11 @@ def get_ssh_client_path(ssh_command="ssh", ssh_client_folder=None):
         if os.path.isfile(ssh_path):
             logger.debug("Attempting to run %s from path %s", ssh_command, ssh_path)
             return ssh_path
-        logger.warning("Could not find %s in provided --ssh-client-folder %s. "
-                       "Attempting to get pre-installed OpenSSH bits.", ssh_command, ssh_client_folder)
+        warn_key = (ssh_command, os.path.abspath(ssh_client_folder))
+        if warn_key not in _warned_ssh_client_folders:
+            logger.warning("Could not find %s in provided --ssh-client-folder %s. "
+                           "Attempting to get pre-installed OpenSSH bits.", ssh_command, ssh_client_folder)
+            _warned_ssh_client_folders.add(warn_key)
 
     if platform.system() != 'Windows':
         return ssh_command

src/sftp/azext_sftp/tests/latest/test_custom.py

@@ -86,31 +86,6 @@ def test_sftp_cert(self, mock_write_cert, mock_get_keys, mock_abspath, mock_isdi
         mock_get_keys.assert_called_once_with('/pubkey/path', None, None, None)
         mock_write_cert.assert_called_once_with(cmd, 'pubkey', '/cert/path', None)
 
-    def test_sftp_connect_preprod(self):
-        """Test SFTP connection to preprod environment.
-        
-        Owner: johnli1
-        """
-        cmd = mock.Mock()
-        cmd.cli_ctx = mock.Mock()
-        cmd.cli_ctx.cloud = mock.Mock()
-        cmd.cli_ctx.cloud.name = "azurecloud"
-        
-        # Create a temporary batch file for automated testing
-        batch_file = os.path.join(self.temp_dir, "test_batch.txt")
-        with open(batch_file, 'w') as f:
-            f.write("pwd\nls\nexit\n")
-        
-        # Use batch file to avoid interactive prompt
-        custom.sftp_connect(
-            cmd=cmd,
-            storage_account='johnli1canary',
-            port=22,
-            cert_file='C:\\Users\\johnli1\\.ssh\\id_rsa-aadcert.pub',
-            sftp_args=['-b', batch_file]  # Use actual batch file
-        )
-        self.assertTrue(True)
-
     @mock.patch('azext_sftp.custom._do_sftp_op')
     @mock.patch('azext_sftp.sftp_utils.get_ssh_cert_principals')
     def test_sftp_connect_certificate_scenarios(self, mock_get_principals, mock_do_sftp):
@@ -713,7 +688,7 @@ def test_sftp_cert_error_cases(self):
         """Test sftp cert error handling with invalid argument combinations."""
         # Test cases: (cert_path, public_key_file, setup_mocks, expected_exception, expected_message, description)
         test_cases = [
-            (None, None, {}, azclierror.RequiredArgumentMissingError, "--output-file or --public-key-file must be provided", "no_arguments"),
+            (None, None, {}, azclierror.RequiredArgumentMissingError, "--file or --public-key-file must be provided", "no_arguments"),
             ("/bad/cert.pub", None, {"expanduser_return": "/bad/cert.pub", "isdir_return": False}, azclierror.InvalidArgumentValueError, "folder doesn't exist", "invalid_directory"),
         ]
 

src/sftp/azext_sftp/tests/latest/test_file_utils.py

@@ -379,11 +379,14 @@ def test_check_or_create_public_private_files_generates_keys(self, mock_create_k
         expected_public_key = os.path.join(self.temp_dir, "id_rsa.pub")
         expected_private_key = os.path.join(self.temp_dir, "id_rsa")
 
-        # Create the expected files so they exist for the function
-        with open(expected_public_key, 'w') as f:
-            f.write("ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ [email protected]")
-        with open(expected_private_key, 'w') as f:
-            f.write("-----BEGIN OPENSSH PRIVATE KEY-----\ntest\n-----END OPENSSH PRIVATE KEY-----")
+        # Mock the create_ssh_keyfile to create the files when called
+        def create_key_files(private_key_path, ssh_client_folder):
+            with open(private_key_path, 'w') as f:
+                f.write("-----BEGIN OPENSSH PRIVATE KEY-----\ntest\n-----END OPENSSH PRIVATE KEY-----")
+            with open(private_key_path + ".pub", 'w') as f:
+                f.write("ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ [email protected]")
+        
+        mock_create_keyfile.side_effect = create_key_files
 
         # Act
         public_key, private_key, delete_keys = file_utils.check_or_create_public_private_files(
@@ -454,7 +457,7 @@ def test_get_and_write_certificate_success(self, mock_set_mode, mock_get_princip
         mock_get_principals.return_value = ["[email protected]"]
 
         # Set up the cert file path that the function will generate
-        expected_cert_file = str(self.mock_public_key) + "-aadcert.pub"
+        expected_cert_file = str(self.mock_public_key.removesuffix(".pub")) + "-aadcert.pub"
         mock_write_cert.return_value = expected_cert_file
 
         # Act

src/sftp/azext_sftp/tests/latest/test_sftp_info.py

@@ -190,6 +190,23 @@ def test_invalid_port_handling(self):
                     # Port should be converted to string
                     self.assertIsInstance(command_args[port_index + 1], str)
 
+    def test_extension_sftp_session_creation(self):
+        """Test that the extension creates SFTP session with correct parameters."""
+        session = sftp_info.SFTPSession(
+            storage_account=self.test_storage_account,
+            username=self.test_username,
+            host=self.test_host,
+            port=self.test_port,
+            cert_file=self.test_cert_file,
+            private_key_file=self.test_private_key_file
+        )
+        
+        self.assertEqual(session.storage_account, self.test_storage_account)
+        self.assertEqual(session.username, self.test_username)
+        self.assertEqual(session.host, self.test_host)
+        self.assertEqual(session.port, self.test_port)
+        self.assertEqual(session.cert_file, os.path.abspath(self.test_cert_file))
+        self.assertEqual(session.private_key_file, os.path.abspath(self.test_private_key_file))
 
 if __name__ == '__main__':
     unittest.main()