fix inconsistence between the get_conn and run in WinRMHook

providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py

@@ -118,13 +118,9 @@ def __init__(
         self.credssp_disable_tlsv1_2 = credssp_disable_tlsv1_2
         self.send_cbt = send_cbt
 
-        self.client = None
         self.winrm_protocol = None
 
     def get_conn(self):
-        if self.client:
-            return self.client
-
         self.log.debug("Creating WinRM client for conn_id: %s", self.ssh_conn_id)
         if self.ssh_conn_id is not None:
             conn = self.get_connection(self.ssh_conn_id)
@@ -213,15 +209,16 @@ def get_conn(self):
                     send_cbt=self.send_cbt,
                 )
 
-            self.log.info("Establishing WinRM connection to host: %s", self.remote_host)
-            self.client = self.winrm_protocol.open_shell()
-
         except Exception as error:
-            error_msg = f"Error connecting to host: {self.remote_host}, error: {error}"
+            error_msg = f"Error creating connection to host: {self.remote_host}, error: {error}"
             self.log.error(error_msg)
             raise AirflowException(error_msg)
 
-        return self.client
+        if not hasattr(self.winrm_protocol, "get_command_output_raw"):
+            # since pywinrm>=0.5 get_command_output_raw replace _raw_get_command_output
+            self.winrm_protocol.get_command_output_raw = self.winrm_protocol._raw_get_command_output
+
+        return self.winrm_protocol
 
     def run(
         self,
@@ -241,18 +238,25 @@ def run(
         :return: returns a tuple containing return_code, stdout and stderr in order.
         """
         winrm_client = self.get_conn()
+        self.log.info("Establishing WinRM connection to host: %s", self.remote_host)
+        try:
+            shell_id = winrm_client.open_shell()
+        except Exception as error:
+            error_msg = f"Error connecting to host: {self.remote_host}, error: {error}"
+            self.log.error(error_msg)
+            raise AirflowException(error_msg)
 
         try:
             if ps_path is not None:
                 self.log.info("Running command as powershell script: '%s'...", command)
                 encoded_ps = b64encode(command.encode("utf_16_le")).decode("ascii")
-                command_id = self.winrm_protocol.run_command(  # type: ignore[attr-defined]
-                    winrm_client, f"{ps_path} -encodedcommand {encoded_ps}"
+                command_id = winrm_client.run_command(  # type: ignore[attr-defined]
+                    shell_id, f"{ps_path} -encodedcommand {encoded_ps}"
                 )
             else:
                 self.log.info("Running command: '%s'...", command)
-                command_id = self.winrm_protocol.run_command(  # type: ignore[attr-defined]
-                    winrm_client, command
+                command_id = winrm_client.run_command(  # type: ignore[attr-defined]
+                    shell_id, command
                 )
 
                 # See: https://github.com/diyan/pywinrm/blob/master/winrm/protocol.py
@@ -267,8 +271,8 @@ def run(
                         stderr,
                         return_code,
                         command_done,
-                    ) = self.winrm_protocol._raw_get_command_output(  # type: ignore[attr-defined]
-                        winrm_client, command_id
+                    ) = winrm_client.get_command_output_raw(  # type: ignore[attr-defined]
+                        shell_id, command_id
                     )
 
                     # Only buffer stdout if we need to so that we minimize memory usage.
@@ -281,12 +285,12 @@ def run(
                     for line in stderr.decode(output_encoding).splitlines():
                         self.log.warning(line)
 
-            self.winrm_protocol.cleanup_command(  # type: ignore[attr-defined]
-                winrm_client, command_id
+            winrm_client.cleanup_command(  # type: ignore[attr-defined]
+                shell_id, command_id
             )
 
             return return_code, stdout_buffer, stderr_buffer
         except Exception as e:
             raise AirflowException(f"WinRM operator error: {e}")
         finally:
-            self.winrm_protocol.close_shell(winrm_client)  # type: ignore[attr-defined]
+            winrm_client.close_shell(shell_id)  # type: ignore[attr-defined]

providers/microsoft/winrm/tests/unit/microsoft/winrm/hooks/test_winrm.py

@@ -29,15 +29,6 @@
 
 
 class TestWinRMHook:
-    @patch("airflow.providers.microsoft.winrm.hooks.winrm.Protocol")
-    def test_get_conn_exists(self, mock_protocol):
-        winrm_hook = WinRMHook()
-        winrm_hook.client = mock_protocol.return_value.open_shell.return_value
-
-        conn = winrm_hook.get_conn()
-
-        assert conn == winrm_hook.client
-
     def test_get_conn_missing_remote_host(self):
         with pytest.raises(AirflowException):
             WinRMHook().get_conn()
@@ -47,7 +38,7 @@ def test_get_conn_error(self, mock_protocol):
         mock_protocol.side_effect = Exception("Error")
 
         with pytest.raises(AirflowException):
-            WinRMHook(remote_host="host").get_conn()
+            WinRMHook(remote_host="host", password="pwd").get_conn()
 
     @patch("airflow.providers.microsoft.winrm.hooks.winrm.Protocol", autospec=True)
     @patch(
@@ -151,7 +142,7 @@ def test_run_with_stdout(self, mock_get_connection, mock_protocol):
         winrm_hook = WinRMHook(ssh_conn_id="conn_id")
 
         mock_protocol.return_value.run_command = MagicMock(return_value="command_id")
-        mock_protocol.return_value._raw_get_command_output = MagicMock(
+        mock_protocol.return_value.get_command_output_raw = MagicMock(
             return_value=(b"stdout", b"stderr", 0, True)
         )
 
@@ -192,7 +183,7 @@ def test_run_without_stdout(self, mock_get_connection, mock_protocol):
         winrm_hook = WinRMHook(ssh_conn_id="conn_id")
 
         mock_protocol.return_value.run_command = MagicMock(return_value="command_id")
-        mock_protocol.return_value._raw_get_command_output = MagicMock(
+        mock_protocol.return_value.get_command_output_raw = MagicMock(
             return_value=(b"stdout", b"stderr", 0, True)
         )