From e3fc0df09e047af675d6b80820dff5cd33dc5cb5 Mon Sep 17 00:00:00 2001
From: Sebastian Villena <97059974+ruisebas@users.noreply.github.com>
Date: Fri, 1 Sep 2023 10:40:45 -0400
Subject: [PATCH] fix(PushNotifications): Adding missing escape hatch to the
 plugin.

---
 ...ntPushNotificationsPlugin+ClientBehaviour.swift |  8 ++++++++
 ...shNotificationsPluginClientBehaviourTests.swift | 14 ++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/AmplifyPlugins/Notifications/Push/Sources/AWSPinpointPushNotificationsPlugin/AWSPinpointPushNotificationsPlugin+ClientBehaviour.swift b/AmplifyPlugins/Notifications/Push/Sources/AWSPinpointPushNotificationsPlugin/AWSPinpointPushNotificationsPlugin+ClientBehaviour.swift
index 1b09110c14..2c615a7cab 100644
--- a/AmplifyPlugins/Notifications/Push/Sources/AWSPinpointPushNotificationsPlugin/AWSPinpointPushNotificationsPlugin+ClientBehaviour.swift
+++ b/AmplifyPlugins/Notifications/Push/Sources/AWSPinpointPushNotificationsPlugin/AWSPinpointPushNotificationsPlugin+ClientBehaviour.swift
@@ -6,6 +6,7 @@
 //
 
 import Amplify
+import AWSPinpoint
 import Foundation
 @_spi(InternalAWSPinpoint) import InternalAWSPinpoint
 import UserNotifications
@@ -59,6 +60,13 @@ extension AWSPinpointPushNotificationsPlugin {
     }
 #endif
 
+    /// Retrieves the escape hatch to perform actions directly on PinpointClient.
+    ///
+    /// - Returns: PinpointClientProtocol instance
+    public func getEscapeHatch() -> PinpointClientProtocol {
+        pinpoint.pinpointClient
+    }
+
     private func recordNotification(_ userInfo: [String: Any],
                                     applicationState: ApplicationState,
                                     action: PushNotification.Action) async {
diff --git a/AmplifyPlugins/Notifications/Push/Tests/AWSPinpointPushNotificationsPluginUnitTests/AWSPinpointPushNotificationsPluginClientBehaviourTests.swift b/AmplifyPlugins/Notifications/Push/Tests/AWSPinpointPushNotificationsPluginUnitTests/AWSPinpointPushNotificationsPluginClientBehaviourTests.swift
index fdf2bd72f9..f7ed2622d3 100644
--- a/AmplifyPlugins/Notifications/Push/Tests/AWSPinpointPushNotificationsPluginUnitTests/AWSPinpointPushNotificationsPluginClientBehaviourTests.swift
+++ b/AmplifyPlugins/Notifications/Push/Tests/AWSPinpointPushNotificationsPluginUnitTests/AWSPinpointPushNotificationsPluginClientBehaviourTests.swift
@@ -6,6 +6,7 @@
 //
 
 import Amplify
+import AWSPinpoint
 @_spi(InternalAWSPinpoint) @testable import InternalAWSPinpoint
 @testable import AWSPinpointPushNotificationsPlugin
 import UserNotifications
@@ -206,6 +207,19 @@ class AWSPinpointPushNotificationsPluginClientBehaviourTests: AWSPinpointPushNot
     }
 #endif
     
+    // - MARK: Escape Hatch tests
+    /// Given: A configured AWSPinpointPushNotificationsPlugin instance
+    /// When: The getEscapeHatch API is invoked
+    /// Then: The underlying PinpointClientProtocol instance is retrieved
+    func testGetEscapeHatch_shouldReturnPinpointClient() {
+        guard let escapeHatch = plugin.getEscapeHatch() as? PinpointClient else {
+            XCTFail("Unable to retrieve PinpointClient")
+            return
+        }
+        
+        XCTAssertTrue(escapeHatch === (mockPinpoint.pinpointClient as? PinpointClient))
+    }
+    
     private func createUserInfo(for source: PushNotification.Source) -> Notifications.Push.UserInfo {
         return [
             "data": [