update view model unit tests

WooCommerce/src/test/kotlin/com/woocommerce/android/ui/orders/wooshippinglabels/address/origin/WooShippingEditOriginViewModelTest.kt

@@ -1,6 +1,7 @@
 package com.woocommerce.android.ui.orders.wooshippinglabels.address.origin
 
 import androidx.compose.runtime.snapshots.Snapshot
+import com.woocommerce.android.model.Location
 import com.woocommerce.android.ui.orders.wooshippinglabels.address.AddressValidationHelper
 import com.woocommerce.android.ui.orders.wooshippinglabels.models.OriginShippingAddress
 import com.woocommerce.android.viewmodel.BaseUnitTest
@@ -17,13 +18,22 @@ import org.mockito.kotlin.whenever
 @OptIn(ExperimentalCoroutinesApi::class)
 class WooShippingEditOriginViewModelTest : BaseUnitTest() {
     private val addressValidator: AddressValidationHelper = mock()
+    private val getAcceptedOriginCountries: GetAcceptedOriginCountries = mock()
+
+    private val countries = listOf(
+        Location("US", "United States"),
+        Location("UK", "United Kingdom"),
+        Location("AR", "Argentina"),
+        Location("BR", "Brazil")
+    )
 
     private lateinit var sut: WooShippingEditOriginViewModel
 
     fun createViewModel(originAddress: OriginShippingAddress) {
         sut = WooShippingEditOriginViewModel(
             addressValidator = addressValidator,
-            savedState = WooShippingEditOriginAddressFragmentArgs(originAddress).toSavedStateHandle()
+            savedState = WooShippingEditOriginAddressFragmentArgs(originAddress).toSavedStateHandle(),
+            getAcceptedOriginCountries = getAcceptedOriginCountries
         )
     }
 
@@ -352,4 +362,44 @@ class WooShippingEditOriginViewModelTest : BaseUnitTest() {
 
         assertThat(isCompanyExpanded).isFalse()
     }
+
+    @Test
+    fun `when get accepted countries succeed then don't display loading or error`() = testBlocking {
+        val address = OriginShippingAddress.EMPTY
+        whenever(addressValidator.validateFieldRequired(any())).doReturn(null)
+        whenever(getAcceptedOriginCountries.invoke()).doReturn(Result.success(countries))
+        Snapshot.withMutableSnapshot {
+            createViewModel(address)
+        }
+
+        advanceUntilIdle()
+
+        val result = sut.viewState.value
+
+        assertThat(result).isInstanceOf(WooShippingEditOriginViewModel.EditAddressViewState::class.java)
+        val dataState = (result as WooShippingEditOriginViewModel.EditAddressViewState.DataState)
+
+        assertThat(dataState.shouldDisplayLoadingCountries).isFalse()
+        assertThat(dataState.shouldDisplayLoadingCountriesError).isFalse()
+    }
+
+    @Test
+    fun `when get accepted countries fails then display error`() = testBlocking {
+        val address = OriginShippingAddress.EMPTY
+        whenever(addressValidator.validateFieldRequired(any())).doReturn(null)
+        whenever(getAcceptedOriginCountries.invoke()).doReturn(Result.failure(Exception("error")))
+        Snapshot.withMutableSnapshot {
+            createViewModel(address)
+        }
+
+        advanceUntilIdle()
+
+        val result = sut.viewState.value
+
+        assertThat(result).isInstanceOf(WooShippingEditOriginViewModel.EditAddressViewState::class.java)
+        val dataState = (result as WooShippingEditOriginViewModel.EditAddressViewState.DataState)
+
+        assertThat(dataState.shouldDisplayLoadingCountries).isFalse()
+        assertThat(dataState.shouldDisplayLoadingCountriesError).isTrue()
+    }
 }