(airflow) branch main updated: Add PSRP connection type (#34766)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

malthe pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new bcf26733af Add PSRP connection type (#34766)
bcf26733af is described below

commit bcf26733af8eb869142a08ca1146b561bfe299bf
Author: Arjun Anandkumar <10...@users.noreply.github.com>
AuthorDate: Fri Feb 16 15:46:30 2024 +0100

    Add PSRP connection type (#34766)
    
    A new connection type "psrp" has been added along with a connection test
    
    Co-authored-by: Andrey Anshin <An...@taragol.is>
---
 airflow/providers/microsoft/psrp/hooks/psrp.py    | 10 ++++++++++
 airflow/providers/microsoft/psrp/provider.yaml    | 10 +++++++---
 tests/providers/microsoft/psrp/hooks/test_psrp.py |  6 ++++++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py
index 0edb245302..aa292a3325 100644
--- a/airflow/providers/microsoft/psrp/hooks/psrp.py
+++ b/airflow/providers/microsoft/psrp/hooks/psrp.py
@@ -75,6 +75,11 @@ class PsrpHook(BaseHook):
     or by setting this key as the extra fields of your connection.
     """
 
+    conn_name_attr = "psrp_conn_id"
+    default_conn_name = "psrp_default"
+    conn_type = "psrp"
+    hook_name = "PowerShell Remoting Protocol"
+
     _conn: RunspacePool | None = None
     _wsman_ref: WeakKeyDictionary[RunspacePool, WSMan] = WeakKeyDictionary()
 
@@ -281,3 +286,8 @@ class PsrpHook(BaseHook):
             log(INFO, "Progress: %s (%s)", record.activity, record.description)
         else:
             log(WARNING, "Unsupported message type: %s", message_type)
+
+    def test_connection(self):
+        """Test PSRP Connection."""
+        with PsrpHook(psrp_conn_id=self.conn_id):
+            pass
diff --git a/airflow/providers/microsoft/psrp/provider.yaml b/airflow/providers/microsoft/psrp/provider.yaml
index 55a11da1cf..0f295321d7 100644
--- a/airflow/providers/microsoft/psrp/provider.yaml
+++ b/airflow/providers/microsoft/psrp/provider.yaml
@@ -19,9 +19,9 @@
 package-name: apache-airflow-providers-microsoft-psrp
 name: PowerShell Remoting Protocol (PSRP)
 description: |
-    This package provides remote execution capabilities via the
-    `PowerShell Remoting Protocol (PSRP)
-    <https://docs.microsoft.com/openspecs/windows_protocols/ms-psrp/>`__.
+  This package provides remote execution capabilities via the
+  `PowerShell Remoting Protocol (PSRP)
+  <https://docs.microsoft.com/openspecs/windows_protocols/ms-psrp/>`__.
 
 state: ready
 source-date-epoch: 1705912159
@@ -60,3 +60,7 @@ hooks:
   - integration-name: Windows PowerShell Remoting Protocol
     python-modules:
       - airflow.providers.microsoft.psrp.hooks.psrp
+
+connection-types:
+  - hook-class-name: airflow.providers.microsoft.psrp.hooks.psrp.PsrpHook
+    connection-type: psrp
diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py
index de998387ef..72989d8f8c 100644
--- a/tests/providers/microsoft/psrp/hooks/test_psrp.py
+++ b/tests/providers/microsoft/psrp/hooks/test_psrp.py
@@ -210,3 +210,9 @@ class TestPsrpHook:
         hook = PsrpHook(CONNECTION_ID)
         ps = hook.invoke_powershell("foo")
         assert call("foo") in ps.add_script.mock_calls
+
+    def test_test_connection(self, runspace_pool, *mocks):
+        connection = Connection(conn_type="psrp")
+        connection.test_connection()
+
+        assert runspace_pool.return_value.__enter__.mock_calls == [call()]