【Storyboard】遅延初期化しているViewModelをDIしてViewControllerをテスト可能にする

# はじめに

ViewModelのinitializerに引数がある場合に、ViewController側で以下のような遅延初期化をすることが多々あります。
このViewControllerにViewModelを注入してテストをしたい場合に、どうすれば注入することができるようになるでしょうか。

本投稿では以下の2点を重視して、既存のプロジェクトでも導入できる方法を解説していきます。

- Storyboardを用いてViewControllerが初期化される
- ViewModelの既存のインターフェースを大きく変更しない

※ Combine.frameworkを利用していますが、RxSwiftでも同様の実装が応用できます。
※ GitHubで[ソースコード](https://github.com/marty-suzuki/LazyInitializedViewModelInjectionSample)も公開しています。

```swift:ベースとなるViewControllerの実装
final class CounterViewController: UIViewController {
    @IBOutlet weak var incrementButton: UIButton!
    @IBOutlet weak var decrementButton: UIButton!
    @IBOutlet weak var countLabel: UILabel!

    private lazy var viewModel = CounterViewModel( // 👈 このViewModelをどう注入するか
        count: 0,
        increment: incrementButton.extension.tap(),
        decrement: decrementButton.extension.tap()
    )
    private var cancellables: [AnyCancellable] = []

    override func viewDidLoad() {
        super.viewDidLoad()

        viewModel.countText
            .assign(to: \.text, on: countLabel)
            .store(in: &cancellables)

        viewModel.isDecrementEnabled
            .assign(to: \.isEnabled, on: decrementButton)
            .store(in: &cancellables)
    }
}
```

```swift:ベースとなるViewModelの実装
final class CounterViewModel {
    let countText: AnyPublisher<String?, Never>
    let isDecrementEnabled: AnyPublisher<Bool, Never>

    private var cancellables = [AnyCancellable]()

    init(count: Int,
         increment: AnyPublisher<Void, Never>,
         decrement: AnyPublisher<Void, Never>) {
        let store = Store()
        store.count = count

        self.countText = store.$count.map(String.init)
            .eraseToAnyPublisher()

        self.isDecrementEnabled = store.$count.map { $0 > 0 }
            .eraseToAnyPublisher()

        increment.map { _ in 1 }
            .merge(with: decrement.map { _ in -1 })
            .flatMap { Just(store.count + $0) }
            .assign(to: \.count, on: store)
            .store(in: &cancellables)
    }

    private final class Store {
        @Published var count = 0
    }
}
```

# ViewModelのInitializeをClosureとして注入する方法

まずは少しの変更で対応できる方法を解説します。

## 1. CounterViewModelのprotocolを定義する

まずはCounterViewModelのprotocolを定義して、CounterViewModelに採用します。

```swift
protocol CounterViewModelType {
    var countText: AnyPublisher<String?, Never> { get }
    var isDecrementEnabled: AnyPublisher<Bool, Never> { get }
}

final class CounterViewModel: CounterViewModelType {
    let countText: AnyPublisher<String?, Never>
    let isDecrementEnabled: AnyPublisher<Bool, Never>
}
```

## 2. InitializerをClosureとして定義する

InitializerをClosureとして定義すると

```swift
let initViewModel = CounterViewModel.init
let viewModel = initViewModel(0,
                              incrementButton.extension.tap(),
                              decrementButton.extension.tap())
```

という実装ができます。
この実装に対して型を明記すると

```swift
let initViewModel: (Int, AnyPublisher<Void, Never>, AnyPublisher<Void, Never>) -> CounterViewModel = CounterViewModel.init
let viewModel: CounterViewModel = initViewModel(0,
                                                incrementButton.extension.tap(),
                                                decrementButton.extension.tap())
```

となります。
このClosureをtypealiasで定義すると以下のようになります。

```swift
typealias InitViewModel = (Int, AnyPublisher<Void, Never>, AnyPublisher<Void, Never>) -> CounterViewModel

let initViewModel: InitViewModel = CounterViewModel.init
```

しかし、このままでは`CounterViewModel`しか初期化することができないため、返り値が`CounterViewModelType`となるClosureをtypealiasとして定義します。

```swift
typealias InitViewModel = (Int, AnyPublisher<Void, Never>, AnyPublisher<Void, Never>) -> CounterViewModelType

let initViewModel: InitViewModel = CounterViewModel.init
let viewModel: CounterViewModelType = initViewModel(0,
                                                    incrementButton.extension.tap(),
                                                    decrementButton.extension.tap())
```

つまり`CounterViewModelType`が返り値となるため、以下のようにモック化したオブジェクトも扱えるようになります。

```swift
final class MockCounterViewModel: CounterViewModelType {
    var countText: AnyPublisher<String?, Never> { _countText.eraseToAnyPublisher() }
    let _countText = PassthroughSubject<String?, Never>()
    var isDecrementEnabled: AnyPublisher<Bool, Never> { _isDecrementEnabled.eraseToAnyPublisher() }
    let _isDecrementEnabled = PassthroughSubject<Bool, Never>()
}

let initViewModel: InitViewModel = { _, _, _ in MockCounterViewModel() }
let viewModel: CounterViewModelType = initViewModel(0,
                                                    incrementButton.extension.tap(),
                                                    decrementButton.extension.tap())
```

## 3. ViewControllerでInitializerのtypealiasを利用する

InitializerのClosureを`initViewModel`という命名でViewControllerのpropertyとして定義します。
Storyboardからの初期化を想定しているため、Implicitly Unwrapped Optionalで定義しています。
遅延初期化をしているViewModelを、`CounterViewModel.init`の代わりに`initViewModel`を利用して遅延初期化します。

```swift
final class CounterViewController: UIViewController {
    @IBOutlet weak var incrementButton: UIButton!
    @IBOutlet weak var decrementButton: UIButton!

    typealias InitViewModel = (Int, AnyPublisher<Void, Never>, AnyPublisher<Void, Never>) -> CounterViewModelType
    private var initViewModel: InitViewModel!
    private var count: Int!
    private lazy var viewModel: CounterViewModelType = initViewModel(
        count: count,
        increment: incrementButton.extension.tap(),
        decrement: decrementButton.extension.tap()
    )
    
    ...
}
```

StoryboardからViewControllerを初期化して`initViewModel`に任意のInitializerを渡す実装以下のようになります。
このstatic methodからViewControllerの生成を行うことで、実装の順番的にviewDidLoadよりも先にpropertyの再代入が行われることが保証できているので、クラッシュすることはありません。

```swift
extension CounterViewController {
    static func makeFromStoryboard(count: Int, initViewModel: @escaping InitViewModel) -> CounterViewController {
        let storyboard = UIStoryboard(name: "CounterViewController", bundle: nil)
        let viewController = storyboard.instantiateInitialViewController() as! CounterViewController
        viewController.count = count
        viewController.initViewModel = initViewModel
        return viewController
    }
}
```

そしてViewControllerを以下のように初期化することで、CounterViewModelの振る舞いをするViewControllerが利用できます。

```swift
let viewController = CounterViewController.makeFromStoryboard(count: 0, initViewModel: CounterViewModel.init)
```

## テストの実装

インスタンス化したモックオブジェクトをInitializerのClosureで返すようにすることで、ViewControllerの振る舞いをテストすることができます。

```swift
class CounterViewControllerTests: XCTestCase {
    private var viewController: CounterViewController!
    private var viewModel: MockCounterViewModel!

    override func setUp() {
        let viewModel = MockCounterViewModel()
        self.viewModel = viewModel
        self.viewController = CounterViewController.makeFromStoryboard(count: 0) { _, _, _ in
            return viewModel
        }
        self.viewController.loadViewIfNeeded()
    }

    func test_countLabel_text() {
        let expected = UUID().uuidString
        viewModel._countText.send(expected)
        XCTAssertEqual(viewController.countLabel.text, expected)
    }
}

extension CounterViewControllerTests {
    private final class MockCounterViewModel: CounterViewModelType {
        var countText: AnyPublisher<String?, Never> { _countText.eraseToAnyPublisher() }
        let _countText = PassthroughSubject<String?, Never>()
        ...
    }
}
```

# ViewModelのFactoryから注入する方法

[ViewModelのInitializeをClosureとして注入する方法](#viewmodelのinitializeをclosureとして注入する方法)を、より厳密に定義します。

## 1. ViewModelFactoryを定義する

`initViewModel: InitViewModel`の代わりに`factory: ViewModelFactory`からViewModelを初期化できるようにします。

```swift
struct CounterViewModelFactory {
    private let count: Int

    init(count: Int) {
        self.count = count
    }

    func initialize(increment: AnyPublisher<Void, Never>,
                    decrement: AnyPublisher<Void, Never>) -> CounterViewModelType {
        CounterViewModel(count: self.count,
                         increment:increment,
                         decrement:decrement)
    }
}
```

ViewController内で遅延初期化が必要になっている引数は、`increment`と`decrement`なので、`count`は外部から受け取りできる状態にします。
このように実装すると以下のようにViewModelを初期化することができます。

```swift
let viewModel: CounterViewModelType = CounterViewModelFactory(count: 0)
    .initialize(increment: incrementButton.extension.tap(),
                decrement: decrementButton.extension.tap())
```

しかし、この実装で`CounterViewModelType`が返り値となっていますが、内部実装は`CounterViewModel`に固定されてしまっているため、実際は振る舞いを変えることができません。

## 2. ViewModelFactoryのprotocolを定義する

ViewModelFactoryの振る舞いを変更可能にするために、ViewModelFactoryのprotocolを定義します。

```swift
protocol ViewModelFactoryType {
    associatedtype Dependency
    associatedtype ViewModel
    func initialize(_ dependency: Dependency) -> ViewModel
}
```

初期化時の引数を汎用的にするために`Dependency`を定義しているので、`CounterViewModelFactory`に適用します。

```swift
struct CounterViewModelFactory: ViewModelFactoryType {
    ...
    func initialize(_ dependency: (
        increment: AnyPublisher<Void, Never>,
        decrement: AnyPublisher<Void, Never>)
    ) -> CounterViewModelType {
        CounterViewModel(count: self.count,
                         increment: dependency.increment,
                         decrement: dependency.decrement)
    }
}
```

しかし、ViewModelFactoryTypeはassociatedTypeを定義しているため、`let viewModelFactory: ViewModelFactoryType`のような定義ができません。
そこで、`AnyViewModelFactory`を定義します。

```swift
struct AnyViewModelFactory<Dependency, ViewModel>: ViewModelFactoryType {

    private let _initialize: (Dependency) -> ViewModel

    init<Factory: ViewModelFactoryType>(_ factory: Factory) where Factory.Dependency == Dependency, Factory.ViewModel == ViewModel {
        self._initialize = { factory.initialize($0) }
    }

    func initialize(_ dependency: Dependency) -> ViewModel {
        _initialize(dependency)
    }
}
```

AnyViewModelFactoryを用いてViewModelの初期化をすると、以下のようになります。

```swift
let factory = AnyViewModelFactory<CounterViewModelFactory.Dependency, CounterViewModelType>(CounterViewModelFactory(count: 0))
let viewModel: CounterViewModelType = factory
    .initialize(increment: incrementButton.extension.tap(),
                decrement: decrementButton.extension.tap())
```

また、MockCounterViewModelFactoryを定義してAnyViewModelFactoryを用いた場合は以下のようになります。

```swift
struct MockCounterViewModelFactory: ViewModelFactoryType {

    let viewModel: MockCounterViewModel
    
    func initialize(_ dependency: (
        increment: AnyPublisher<Void, Never>,
        decrement: AnyPublisher<Void, Never>)
    ) -> CounterViewModelType {
        viewModel
    }
}

let factory = AnyViewModelFactory<CounterViewModelFactory.Dependency, CounterViewModelType>(MockCounterViewModelFactory(viewModel: MockCounterViewModel()))
let viewModel: CounterViewModelType = factory
    .initialize(increment: incrementButton.extension.tap(),
                decrement: decrementButton.extension.tap())
```

Factoryの型自体は`AnyViewModelFactory<CounterViewModelFactory.Dependency, CounterViewModelType>`になっているため、propertyとして定義して異なる振る舞いを注入することができます。

## 3. ViewControllerでViewModelFactoryを利用する

`factory: AnyViewModelFactory<CounterViewModelFactory.Dependency, CounterViewModelType>`をViewControllerに定義します。
Storyboardからの初期化を想定しているため、Implicitly Unwrapped Optionalで定義しています。
遅延初期化をしているViewModelを、`CounterViewModel.init`の代わりに`factory.initialize`を利用して遅延初期化します。

```swift
final class CounterViewController: UIViewController {
    @IBOutlet weak var incrementButton: UIButton!
    @IBOutlet weak var decrementButton: UIButton!

    private var factory: AnyViewModelFactory<CounterViewModelFactory.Dependency, CounterViewModelType>!
    private lazy var viewModel = factory.initialize((
        increment: incrementButton.extension.tap(),
        decrement: decrementButton.extension.tap()
    ))
```

StoryboardからViewControllerを初期化して`factory`に任意のViewModelFactoryを渡す実装以下のようになります。
Generic ArgumentにViewModelFactoryを指定し、where句で`CounterViewModelFactory.Dependency`と`CounterViewModelType`を持つViewModelFactoryに限定します。
そのViewModelFactoryをAnyViewModelFactory<CounterViewModelFactory.Dependency, CounterViewModelType>で型消去をして、factoryに代入します。
このstatic methodからViewControllerの生成を行うことで、実装の順番的にviewDidLoadよりも先にpropertyの再代入が行われることが保証できているので、クラッシュすることはありません。

```swift
extension CounterViewController {
    static func makeFromStoryboard<Factory: ViewModelFactory>(factory: Factory) -> CounterViewController where Factory.Dependency == CounterViewModelFactory.Dependency, Factory.ViewModel == CounterViewModelType {
        let storyboard = UIStoryboard(name: "CounterViewController", bundle: nil)
        let viewController = storyboard.instantiateInitialViewController() as! CounterViewController
        viewController.factory = AnyViewModelFactory(factory)
        return viewController
    }
}
```

そしてViewControllerを以下のように初期化することで、CounterViewModelの振る舞いをするViewControllerが利用できます。

```swift
let viewController = CounterViewController.makeFromStoryboard(factory: CounterViewModelFactory(count: 0))
```

## テストの実装

モック化したFactoryでモック化したViewModelのインスタンスを返すようにすることで、ViewControllerの振る舞いをテストすることができます。

```swift
class CounterViewControllerTests: XCTestCase {
    private var viewController: CounterViewController!
    private var viewModel: MockCounterViewModel!
    private var factory: MockCounterViewModelFactory!

    override func setUp() {
        self.factory = MockCounterViewModelFactory()
        self.viewModel = factory.viewModel
        self.viewController = CounterViewController.makeFromStoryboard(factory: self.factory)
        viewController.loadViewIfNeeded()
    }

    func test_countLabel_text() {
        let expected = UUID().uuidString
        viewModel._countText.send(expected)
        XCTAssertEqual(viewController.countLabel.text, expected)
    }
}

extension CounterViewControllerTests {
    private final class MockCounterViewModel: CounterViewModelType {
        var countText: AnyPublisher<String?, Never> { _countText.eraseToAnyPublisher() }
        let _countText = PassthroughSubject<String?, Never>()
        ...
    }

    private struct MockCounterViewModelFactory: ViewModelFactoryType {
        let viewModel = MockCounterViewModel()

        func initialize(_ dependency: CounterViewModelFactory.Dependency) -> CounterViewModelType {
            return viewModel
        }
    }
}
```

# 最後に

Storyboardを想定しているため`Implicitly Unwrapped Optional`としていますが、その部分以外は`xib`などでも同じ実装が応用できます。
既存のプロジェクトにも比較的導入しやすい方法であると思うので、是非試してみてください。