Save country when creating new quote addresses

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());

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());

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);

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);