From 00b8df33d71b18a4c97b9616e9442d0ee809ed65 Mon Sep 17 00:00:00 2001
From: Samer Alabi <samer@stripe.com>
Date: Mon, 30 Sep 2024 16:33:55 -0400
Subject: [PATCH] Return `CustomerEphemeralKey.Available` directly rather than
 have consumers handle it.

---
 .../data/CachedCustomerEphemeralKey.kt        | 11 --------
 .../CustomerSessionElementsSessionManager.kt  | 25 ++++++++++++-------
 ...CustomerSessionSavedSelectionDataSource.kt |  4 +--
 ...keCustomerSessionElementsSessionManager.kt |  4 +--
 4 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CachedCustomerEphemeralKey.kt b/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CachedCustomerEphemeralKey.kt
index f14e0602fbd..d82f21c991d 100644
--- a/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CachedCustomerEphemeralKey.kt
+++ b/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CachedCustomerEphemeralKey.kt
@@ -23,14 +23,3 @@ internal sealed interface CachedCustomerEphemeralKey {
         }
     }
 }
-
-internal fun <T> Result<CachedCustomerEphemeralKey>.mapWithEphemeralKeyCatching(
-    mapper: (customerId: String, ephemeralKey: String) -> T
-): Result<T> {
-    return mapCatching { cachedKey ->
-        when (cachedKey) {
-            is CachedCustomerEphemeralKey.Available -> mapper(cachedKey.customerId, cachedKey.ephemeralKey)
-            is CachedCustomerEphemeralKey.None -> throw IllegalStateException("No ephemeral key was available!")
-        }
-    }
-}
diff --git a/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionElementsSessionManager.kt b/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionElementsSessionManager.kt
index 4137b71134b..0bc24276373 100644
--- a/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionElementsSessionManager.kt
+++ b/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionElementsSessionManager.kt
@@ -15,7 +15,7 @@ import javax.inject.Singleton
 import kotlin.coroutines.CoroutineContext
 
 internal interface CustomerSessionElementsSessionManager {
-    suspend fun fetchCustomerSessionEphemeralKey(): Result<CachedCustomerEphemeralKey>
+    suspend fun fetchCustomerSessionEphemeralKey(): Result<CachedCustomerEphemeralKey.Available>
 
     suspend fun fetchElementsSession(): Result<ElementsSession>
 }
@@ -34,16 +34,23 @@ internal class DefaultCustomerSessionElementsSessionManager @Inject constructor(
 
     private var intentConfiguration: CustomerSheet.IntentConfiguration? = null
 
-    override suspend fun fetchCustomerSessionEphemeralKey(): Result<CachedCustomerEphemeralKey> {
+    override suspend fun fetchCustomerSessionEphemeralKey(): Result<CachedCustomerEphemeralKey.Available> {
         return withContext(workContext) {
-            cachedCustomerEphemeralKey.takeUnless { cachedCustomerEphemeralKey ->
-                cachedCustomerEphemeralKey.shouldRefresh(timeProvider())
-            }?.let {
-                Result.success(it)
-            } ?: kotlin.runCatching {
-                fetchElementsSession().getOrThrow()
+            runCatching {
+                val ephemeralKey = cachedCustomerEphemeralKey.takeUnless { cachedCustomerEphemeralKey ->
+                    cachedCustomerEphemeralKey.shouldRefresh(timeProvider())
+                } ?: run {
+                    fetchElementsSession().getOrThrow()
+
+                    cachedCustomerEphemeralKey
+                }
 
-                cachedCustomerEphemeralKey
+                when (ephemeralKey) {
+                    is CachedCustomerEphemeralKey.Available -> ephemeralKey
+                    is CachedCustomerEphemeralKey.None -> throw IllegalStateException(
+                        "No ephemeral key available!"
+                    )
+                }
             }
         }
     }
diff --git a/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionSavedSelectionDataSource.kt b/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionSavedSelectionDataSource.kt
index f379fc524ac..04656cadb13 100644
--- a/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionSavedSelectionDataSource.kt
+++ b/paymentsheet/src/main/java/com/stripe/android/customersheet/data/CustomerSessionSavedSelectionDataSource.kt
@@ -42,8 +42,8 @@ internal class CustomerSessionSavedSelectionDataSource @Inject constructor(
     }
 
     private suspend fun createPrefsRepository(): CustomerSheetDataResult<PrefsRepository> {
-        return elementsSessionManager.fetchCustomerSessionEphemeralKey().mapWithEphemeralKeyCatching { customerId, _ ->
-            prefsRepositoryFactory(customerId)
+        return elementsSessionManager.fetchCustomerSessionEphemeralKey().mapCatching { ephemeralKey ->
+            prefsRepositoryFactory(ephemeralKey.customerId)
         }.toCustomerSheetDataResult()
     }
 }
diff --git a/paymentsheet/src/test/java/com/stripe/android/customersheet/data/FakeCustomerSessionElementsSessionManager.kt b/paymentsheet/src/test/java/com/stripe/android/customersheet/data/FakeCustomerSessionElementsSessionManager.kt
index 0a3d9a227f0..3a9f46a798a 100644
--- a/paymentsheet/src/test/java/com/stripe/android/customersheet/data/FakeCustomerSessionElementsSessionManager.kt
+++ b/paymentsheet/src/test/java/com/stripe/android/customersheet/data/FakeCustomerSessionElementsSessionManager.kt
@@ -3,7 +3,7 @@ package com.stripe.android.customersheet.data
 import com.stripe.android.model.ElementsSession
 
 internal class FakeCustomerSessionElementsSessionManager(
-    private val ephemeralKey: Result<CachedCustomerEphemeralKey> = Result.success(
+    private val ephemeralKey: Result<CachedCustomerEphemeralKey.Available> = Result.success(
         CachedCustomerEphemeralKey.Available(
             customerId = "cus_1",
             ephemeralKey = "ek_123",
@@ -11,7 +11,7 @@ internal class FakeCustomerSessionElementsSessionManager(
         )
     ),
 ) : CustomerSessionElementsSessionManager {
-    override suspend fun fetchCustomerSessionEphemeralKey(): Result<CachedCustomerEphemeralKey> {
+    override suspend fun fetchCustomerSessionEphemeralKey(): Result<CachedCustomerEphemeralKey.Available> {
         return ephemeralKey
     }