From e69b74e5a2218b39c77e97a321e1133361e495e7 Mon Sep 17 00:00:00 2001
From: Olivier Daneau <olivier.daneau@astronomer.io>
Date: Fri, 25 Nov 2022 15:07:52 -0500
Subject: [PATCH] Add test_connection to SalesforceHook

---
 .../providers/salesforce/hooks/salesforce.py  | 12 +++++++++++
 .../salesforce/hooks/test_salesforce.py       | 20 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/airflow/providers/salesforce/hooks/salesforce.py b/airflow/providers/salesforce/hooks/salesforce.py
index 8087c42b07187..6d2ccd50277c4 100644
--- a/airflow/providers/salesforce/hooks/salesforce.py
+++ b/airflow/providers/salesforce/hooks/salesforce.py
@@ -391,3 +391,15 @@ def object_to_df(
             df["time_fetched_from_salesforce"] = fetched_time
 
         return df
+
+    def test_connection(self):
+        """Test the Salesforce connectivity"""
+        try:
+            self.describe_object("Account")
+            status = True
+            message = "Connection successfully tested"
+        except Exception as e:
+            status = False
+            message = str(e)
+
+        return status, message
diff --git a/tests/providers/salesforce/hooks/test_salesforce.py b/tests/providers/salesforce/hooks/test_salesforce.py
index 4b936c32fe1e1..c515a1ee0382f 100644
--- a/tests/providers/salesforce/hooks/test_salesforce.py
+++ b/tests/providers/salesforce/hooks/test_salesforce.py
@@ -514,3 +514,23 @@ def test_backcompat_prefix_both_prefers_short(self, mock_client):
                 username=None,
                 version="52.0",
             )
+
+    @patch(
+        "airflow.providers.salesforce.hooks.salesforce.SalesforceHook.describe_object",
+        return_value={"fields": [{"name": "field_1"}, {"name": "field_2"}]},
+    )
+    def test_connection_success(self, mock_describe_object):
+        hook = SalesforceHook("my_conn")
+        status, msg = hook.test_connection()
+        assert status is True
+        assert msg == "Connection successfully tested"
+
+    @patch(
+        "airflow.providers.salesforce.hooks.salesforce.SalesforceHook.describe_object",
+        side_effect=Exception("Test"),
+    )
+    def test_connection_failure(self, mock_describe_object):
+        hook = SalesforceHook("my_conn")
+        status, msg = hook.test_connection()
+        assert status is False
+        assert msg == "Test"