From 7871f4c973102ef12af1748ff296bc4859aa0049 Mon Sep 17 00:00:00 2001
From: Patrick McLain <pat@pmclain.com>
Date: Fri, 22 Feb 2019 20:26:26 -0500
Subject: [PATCH] Save country when creating new quote addresses

Fixes #389
---
 .../Model/Cart/SetBillingAddressOnCart.php    |  1 +
 .../Model/Cart/SetShippingAddressOnCart.php   |  1 +
 .../Quote/SetBillingAddressOnCartTest.php     | 19 ++++++++++++++++++-
 .../Quote/SetShippingAddressOnCartTest.php    | 11 ++++++++++-
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/app/code/Magento/QuoteGraphQl/Model/Cart/SetBillingAddressOnCart.php b/app/code/Magento/QuoteGraphQl/Model/Cart/SetBillingAddressOnCart.php
index 5d3125ae13ef8..02aec5b6fbaf0 100644
--- a/app/code/Magento/QuoteGraphQl/Model/Cart/SetBillingAddressOnCart.php
+++ b/app/code/Magento/QuoteGraphQl/Model/Cart/SetBillingAddressOnCart.php
@@ -93,6 +93,7 @@ public function execute(ContextInterface $context, CartInterface $cart, array $b
             );
         }
         if (null === $customerAddressId) {
+            $addressInput['country_id'] = $addressInput['country_code'] ?? '';
             $billingAddress = $this->addressModel->addData($addressInput);
         } else {
             $this->checkCustomerAccount->execute($context->getUserId(), $context->getUserType());
diff --git a/app/code/Magento/QuoteGraphQl/Model/Cart/SetShippingAddressOnCart.php b/app/code/Magento/QuoteGraphQl/Model/Cart/SetShippingAddressOnCart.php
index e6b18fc88a27a..fc8eff4cfc13a 100644
--- a/app/code/Magento/QuoteGraphQl/Model/Cart/SetShippingAddressOnCart.php
+++ b/app/code/Magento/QuoteGraphQl/Model/Cart/SetShippingAddressOnCart.php
@@ -87,6 +87,7 @@ public function execute(ContextInterface $context, CartInterface $cart, array $s
             );
         }
         if (null === $customerAddressId) {
+            $addressInput['country_id'] = $addressInput['country_code'] ?? '';
             $shippingAddress = $this->addressModel->addData($addressInput);
         } else {
             $this->checkCustomerAccount->execute($context->getUserId(), $context->getUserType());
diff --git a/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetBillingAddressOnCartTest.php b/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetBillingAddressOnCartTest.php
index bfe57109d107d..754b157a16107 100644
--- a/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetBillingAddressOnCartTest.php
+++ b/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetBillingAddressOnCartTest.php
@@ -88,6 +88,10 @@ public function testSetNewBillingAddressByGuest()
         city
         postcode
         telephone
+        country {
+          code
+          label
+        }
       }
     }
   }
@@ -140,6 +144,10 @@ public function testSetNewBillingAddressWithUseForShippingParameterByGuest()
         city
         postcode
         telephone
+        country {
+          code
+          label
+        }
       }
       shipping_addresses {
         firstname
@@ -149,6 +157,10 @@ public function testSetNewBillingAddressWithUseForShippingParameterByGuest()
         city
         postcode
         telephone
+        country {
+          code
+          label
+        }
       }
     }
   }
@@ -234,6 +246,10 @@ public function testSetNewBillingAddressByRegisteredCustomer()
         city
         postcode
         telephone
+        country {
+          code
+          label
+        }
       }
     }
   }
@@ -413,7 +429,8 @@ private function assertNewAddressFields(array $billingAddressResponse): void
             ['response_field' => 'street', 'expected_value' => [0 => 'test street 1', 1 => 'test street 2']],
             ['response_field' => 'city', 'expected_value' => 'test city'],
             ['response_field' => 'postcode', 'expected_value' => '887766'],
-            ['response_field' => 'telephone', 'expected_value' => '88776655']
+            ['response_field' => 'telephone', 'expected_value' => '88776655'],
+            ['response_field' => 'country', 'expected_value' => ['code' => 'US', 'label' => 'US']],
         ];
 
         $this->assertResponseFields($billingAddressResponse, $assertionMap);
diff --git a/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetShippingAddressOnCartTest.php b/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetShippingAddressOnCartTest.php
index c3685c88ae487..4d7ab85cd475a 100644
--- a/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetShippingAddressOnCartTest.php
+++ b/dev/tests/api-functional/testsuite/Magento/GraphQl/Quote/SetShippingAddressOnCartTest.php
@@ -90,6 +90,10 @@ public function testSetNewShippingAddressByGuest()
         city
         postcode
         telephone
+        country {
+          code
+          label
+        }
       }
     }
   }
@@ -176,6 +180,10 @@ public function testSetNewShippingAddressByRegisteredCustomer()
         city
         postcode
         telephone
+        country {
+          label
+          code
+        }
       }
     }
   }
@@ -462,7 +470,8 @@ private function assertNewShippingAddressFields(array $shippingAddressResponse):
             ['response_field' => 'street', 'expected_value' => [0 => 'test street 1', 1 => 'test street 2']],
             ['response_field' => 'city', 'expected_value' => 'test city'],
             ['response_field' => 'postcode', 'expected_value' => '887766'],
-            ['response_field' => 'telephone', 'expected_value' => '88776655']
+            ['response_field' => 'telephone', 'expected_value' => '88776655'],
+            ['response_field' => 'country', 'expected_value' => ['code' => 'US', 'label' => 'US']],
         ];
 
         $this->assertResponseFields($shippingAddressResponse, $assertionMap);