From 3cb8150c9a7ab8d5d5e3712867290dcdce4c5afa Mon Sep 17 00:00:00 2001
From: Oleg Tverdokhleb <oltv00@gmail.com>
Date: Tue, 16 Oct 2018 12:19:13 +0300
Subject: [PATCH] Fix one payment option method at return

---
 .../Presenter/TokenizationPresenter.swift          | 14 +++++++++++++-
 .../Tokenization/Strategy/ApplePayStrategy.swift   |  2 +-
 .../Strategy/TokenizationStrategy.swift            |  2 ++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/YandexCheckoutPayments/Modules/Tokenization/Presenter/TokenizationPresenter.swift b/YandexCheckoutPayments/Modules/Tokenization/Presenter/TokenizationPresenter.swift
index aa424312..540be821 100644
--- a/YandexCheckoutPayments/Modules/Tokenization/Presenter/TokenizationPresenter.swift
+++ b/YandexCheckoutPayments/Modules/Tokenization/Presenter/TokenizationPresenter.swift
@@ -238,6 +238,18 @@ extension TokenizationPresenter: TokenizationStrategyOutput {
                                            moduleOutput: strongSelf)
         }
     }
+
+    func didFinish(on module: TokenizationStrategyInput) {
+        handleOnePaymentOptionMethodAtReturn()
+    }
+
+    func handleOnePaymentOptionMethodAtReturn() {
+        if paymentOptionsCount == 1 {
+            close()
+        } else {
+            presentPaymentMethodsModule()
+        }
+    }
 }
 
 // MARK: - TokenizationViewOutput
@@ -456,7 +468,7 @@ extension TokenizationPresenter: YandexAuthModuleOutput {
     func didFailFetchYamoneyPaymentMethods(on module: YandexAuthModuleInput) {}
 
     func didCancelAuthorizeInYandex(on module: YandexAuthModuleInput) {
-        presentPaymentMethodsModule()
+        handleOnePaymentOptionMethodAtReturn()
     }
 }
 
diff --git a/YandexCheckoutPayments/Modules/Tokenization/Strategy/ApplePayStrategy.swift b/YandexCheckoutPayments/Modules/Tokenization/Strategy/ApplePayStrategy.swift
index a3458796..5b10c5aa 100644
--- a/YandexCheckoutPayments/Modules/Tokenization/Strategy/ApplePayStrategy.swift
+++ b/YandexCheckoutPayments/Modules/Tokenization/Strategy/ApplePayStrategy.swift
@@ -53,7 +53,7 @@ extension ApplePayStrategy: TokenizationStrategyInput {
 
     func paymentAuthorizationViewControllerDidFinish(_ controller: PKPaymentAuthorizationViewController) {
         if case .failed = paymentResult {
-            output?.presentPaymentMethodsModule()
+            output?.didFinish(on: self)
         }
     }
 
diff --git a/YandexCheckoutPayments/Modules/Tokenization/Strategy/TokenizationStrategy.swift b/YandexCheckoutPayments/Modules/Tokenization/Strategy/TokenizationStrategy.swift
index 30e62d30..a2a6c703 100644
--- a/YandexCheckoutPayments/Modules/Tokenization/Strategy/TokenizationStrategy.swift
+++ b/YandexCheckoutPayments/Modules/Tokenization/Strategy/TokenizationStrategy.swift
@@ -75,6 +75,8 @@ protocol TokenizationStrategyOutput: class {
 
     func presentErrorWithMessage(_ message: String)
 
+    func didFinish(on module: TokenizationStrategyInput)
+
     // MARK: - ApplePay
 
     func presentApplePay()