Repository: saleor/saleor
Pull request: Set checkout country to one from shipping address
This also checks if variants in checkouts are available in the shipping zone from shipping country

I want to merge this change because...
Fixes #5780
Fixes #5781 

<!-- Please mention all relevant issue numbers. -->

# Impact

* [ ] New migrations
* [ ] New/Updated API fields or mutations
* [ ] Deprecated API fields or mutations
* [ ] Removed API types, fields, or mutations
* [ ] Documentation needs to be updated

# Pull Request Checklist

<!-- Please keep this section. It will make the maintainer's life easier. -->

* [ ] Privileged queries and mutations are guarded by proper permission checks
* [ ] Database queries are optimized and the number of queries is constant
* [ ] Database migration files are up to date
* [ ] The changes are tested
* [ ] GraphQL schema and type definitions are up to date
* [ ] Changes are mentioned in the changelog


Problem statement:
Can not buy product if product unavailable for the default country
### What I'm trying to achieve
Buy product with is available in the shipping zone.

### Steps to reproduce the problem
1. Create a product with is unavailable in the shipping zone related to `DEFAULT_COUNTRY` and is available in another shipping zone. 
2. Try to create a checkout with an address in another shipping zone.

#### Example
Product is available in Europe, `DEFAULT_COUNTRY` = "US". 
If I try to buy this product with shipping to Poland I got the error:
`{"checkoutCreate": {"errors": [{"code": "INSUFFICIENT_STOCK", "field": "quantity", "message": "Could not add item AAATestEU (AAATestEU). Only 0 remaining in stock.", "__typename": "CheckoutError"}], "checkout": null, "__typename": "CheckoutCreate"}}}]`



Lack of availability validation on `checkoutShippingAddressUpdate` 
### What I'm trying to achieve
Change checkout address. 

### Steps to reproduce the problem
1. Create product available only in the shipping zone related to `DEFAULT_COUNTRY` 
2. Create checkout for any address in `DEFAULT_COUNTRY` 
3. Execute `checkoutShippingAddressUpdate` with an address out of the shipping zone related to `DEFAULT_COUNTRY` 

### What I expected to happen
Validation shouldn't allow changing checkout shipping addresses if the product is unavailable for shipping zone related to a new address.

Hints from issue discussion:

Commit message: Remove return tuple in process_checkout_lines

Patch to review:
diff --git a/saleor/graphql/checkout/mutations.py b/saleor/graphql/checkout/mutations.py
index 61f9a1e39d4..d5f3d172f1d 100644
--- a/saleor/graphql/checkout/mutations.py
+++ b/saleor/graphql/checkout/mutations.py
@@ -225,6 +225,13 @@ def clean_input(cls, info, instance: models.Checkout, data, input_cls=None):
         user = info.context.user
         country = info.context.country.code
 
+        # set country to one from shipping address
+        shipping_address = cleaned_input.get("shipping_address")
+        if shipping_address and shipping_address.country:
+            if shipping_address.country != country:
+                country = shipping_address.country
+        cleaned_input["country"] = country
+
         # Resolve and process the lines, retrieving the variants and quantities
         lines = data.pop("lines", None)
         if lines:
@@ -267,8 +274,8 @@ def save_addresses(cls, instance: models.Checkout, cleaned_input: dict):
     def save(cls, info, instance: models.Checkout, cleaned_input):
         # Create the checkout object
         instance.save()
-        country = info.context.country
-        instance.set_country(country.code, commit=True)
+        country = cleaned_input["country"]
+        instance.set_country(country, commit=True)
 
         # Retrieve the lines to create
         variants = cleaned_input.get("variants")
@@ -507,6 +514,18 @@ class Meta:
         error_type_class = CheckoutError
         error_type_field = "checkout_errors"
 
+    @classmethod
+    def process_checkout_lines(cls, lines, country) -> None:
+        variant_ids = [line.variant.id for line in lines]
+        variants = list(
+            product_models.ProductVariant.objects.filter(
+                id__in=variant_ids
+            ).prefetch_related("product__product_type")
+        )
+        quantities = [line.quantity for line in lines]
+
+        check_lines_quantity(variants, quantities, country)
+
     @classmethod
     def perform_mutation(cls, _root, info, checkout_id, shipping_address):
         pk = from_global_id_strict_type(checkout_id, Checkout, field="checkout_id")
@@ -541,6 +560,17 @@ def perform_mutation(cls, _root, info, checkout_id, shipping_address):
 
         lines = list(checkout)
 
+        country = info.context.country.code
+        # set country to one from shipping address
+        if shipping_address and shipping_address.country:
+            if shipping_address.country != country:
+                country = shipping_address.country
+        checkout.set_country(country, commit=True)
+
+        # Resolve and process the lines, validating variants quantities
+        if lines:
+            cls.process_checkout_lines(lines, country)
+
         update_checkout_shipping_method_if_invalid(
             checkout, lines, info.context.discounts
         )
diff --git a/saleor/graphql/checkout/tests/test_checkout.py b/saleor/graphql/checkout/tests/test_checkout.py
index d330329c195..34e814b7a1d 100644
--- a/saleor/graphql/checkout/tests/test_checkout.py
+++ b/saleor/graphql/checkout/tests/test_checkout.py
@@ -12,7 +12,7 @@
 from ....checkout import calculations
 from ....checkout.error_codes import CheckoutErrorCode
 from ....checkout.models import Checkout
-from ....checkout.utils import is_fully_paid
+from ....checkout.utils import add_variant_to_checkout, is_fully_paid
 from ....core.payments import PaymentInterface
 from ....core.taxes import zero_money
 from ....order.models import Order
@@ -527,6 +527,68 @@ def test_checkout_create_check_lines_quantity_multiple_warehouse(
     assert data["errors"][0]["field"] == "quantity"
 
 
+def test_checkout_create_sets_country_from_shipping_address_country(
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    variant_id = graphene.Node.to_global_id("ProductVariant", variant.id)
+    test_email = "test@example.com"
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = 10001
+    variables = {
+        "checkoutInput": {
+            "lines": [{"quantity": 1, "variantId": variant_id}],
+            "email": test_email,
+            "shippingAddress": shipping_address,
+        }
+    }
+    assert not Checkout.objects.exists()
+    response = user_api_client.post_graphql(MUTATION_CHECKOUT_CREATE, variables)
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutCreate"]
+    assert data["created"] is True
+    checkout = Checkout.objects.first()
+    assert checkout.country == "US"
+
+
+def test_checkout_create_check_lines_quantity_different_shipping_zone_insufficient_stocks(  # NOQA
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    stock = Stock.objects.filter(
+        warehouse__shipping_zones__countries__contains="US", product_variant=variant
+    ).first()
+    stock.quantity = 0
+    stock.save()
+    variant_id = graphene.Node.to_global_id("ProductVariant", variant.id)
+    test_email = "test@example.com"
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = 10001
+    variables = {
+        "checkoutInput": {
+            "lines": [{"quantity": 1, "variantId": variant_id}],
+            "email": test_email,
+            "shippingAddress": shipping_address,
+        }
+    }
+    assert not Checkout.objects.exists()
+    response = user_api_client.post_graphql(MUTATION_CHECKOUT_CREATE, variables)
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutCreate"]
+    assert not data["checkout"]
+    errors = data["checkoutErrors"]
+    assert errors[0]["code"] == CheckoutErrorCode.INSUFFICIENT_STOCK.name
+    assert errors[0]["field"] == "quantity"
+
+
 def test_checkout_create_check_lines_quantity(
     user_api_client, stock, graphql_address_data
 ):
@@ -1278,6 +1340,92 @@ def test_checkout_shipping_address_update(
     )
 
 
+@mock.patch(
+    "saleor.graphql.checkout.mutations.update_checkout_shipping_method_if_invalid",
+    wraps=update_checkout_shipping_method_if_invalid,
+)
+def test_checkout_shipping_address_update_changes_checkout_country(
+    mocked_update_shipping_method,
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    checkout = Checkout.objects.create()
+    checkout.set_country("PL", commit=True)
+    add_variant_to_checkout(checkout, variant, 1)
+    assert checkout.shipping_address is None
+    checkout_id = graphene.Node.to_global_id("Checkout", checkout.pk)
+
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = "10001"
+    variables = {"checkoutId": checkout_id, "shippingAddress": shipping_address}
+
+    response = user_api_client.post_graphql(
+        MUTATION_CHECKOUT_SHIPPING_ADDRESS_UPDATE, variables
+    )
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutShippingAddressUpdate"]
+    assert not data["errors"]
+    checkout.refresh_from_db()
+    assert checkout.shipping_address is not None
+    assert checkout.shipping_address.first_name == shipping_address["firstName"]
+    assert checkout.shipping_address.last_name == shipping_address["lastName"]
+    assert (
+        checkout.shipping_address.street_address_1 == shipping_address["streetAddress1"]
+    )
+    assert (
+        checkout.shipping_address.street_address_2 == shipping_address["streetAddress2"]
+    )
+    assert checkout.shipping_address.postal_code == shipping_address["postalCode"]
+    assert checkout.shipping_address.country == shipping_address["country"]
+    assert checkout.shipping_address.city == shipping_address["city"].upper()
+    mocked_update_shipping_method.assert_called_once_with(
+        checkout, list(checkout), mock.ANY
+    )
+    assert checkout.country == shipping_address["country"]
+
+
+@mock.patch(
+    "saleor.graphql.checkout.mutations.update_checkout_shipping_method_if_invalid",
+    wraps=update_checkout_shipping_method_if_invalid,
+)
+def test_checkout_shipping_address_update_insufficient_stocks(
+    mocked_update_shipping_method,
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    checkout = Checkout.objects.create()
+    checkout.set_country("PL", commit=True)
+    add_variant_to_checkout(checkout, variant, 1)
+    stock = Stock.objects.filter(
+        warehouse__shipping_zones__countries__contains="US", product_variant=variant
+    ).first()
+    stock.quantity = 0
+    stock.save()
+    assert checkout.shipping_address is None
+    checkout_id = graphene.Node.to_global_id("Checkout", checkout.pk)
+
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = "10001"
+    variables = {"checkoutId": checkout_id, "shippingAddress": shipping_address}
+
+    response = user_api_client.post_graphql(
+        MUTATION_CHECKOUT_SHIPPING_ADDRESS_UPDATE, variables
+    )
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutShippingAddressUpdate"]
+    errors = data["checkoutErrors"]
+    assert errors[0]["code"] == CheckoutErrorCode.INSUFFICIENT_STOCK.name
+    assert errors[0]["field"] == "quantity"
+
+
 def test_checkout_shipping_address_with_invalid_phone_number_returns_error(
     user_api_client, checkout_with_item, graphql_address_data
 ):
diff --git a/saleor/tests/fixtures.py b/saleor/tests/fixtures.py
index 1cf9ffd2159..e22c0606e0f 100644
--- a/saleor/tests/fixtures.py
+++ b/saleor/tests/fixtures.py
@@ -495,6 +495,27 @@ def shipping_zone(db):  # pylint: disable=W0613
     return shipping_zone
 
 
+@pytest.fixture
+def shipping_zones(db):
+    shipping_zone_poland = ShippingZone.objects.create(name="Poland", countries=["PL"])
+    shipping_zone_poland.shipping_methods.create(
+        name="DHL",
+        minimum_order_price=Money(0, "USD"),
+        type=ShippingMethodType.PRICE_BASED,
+        price=Money(10, "USD"),
+        shipping_zone=shipping_zone,
+    )
+    shipping_zone_usa = ShippingZone.objects.create(name="USA", countries=["US"])
+    shipping_zone_usa.shipping_methods.create(
+        name="DHL",
+        minimum_order_price=Money(0, "USD"),
+        type=ShippingMethodType.PRICE_BASED,
+        price=Money(10, "USD"),
+        shipping_zone=shipping_zone,
+    )
+    return [shipping_zone_poland, shipping_zone_usa]
+
+
 @pytest.fixture
 def shipping_zone_without_countries(db):  # pylint: disable=W0613
     shipping_zone = ShippingZone.objects.create(name="Europe", countries=[])
@@ -859,6 +880,16 @@ def variant_with_many_stocks(variant, warehouses_with_shipping_zone):
     return variant
 
 
+@pytest.fixture
+def variant_with_many_stocks_different_shipping_zones(
+    variant, warehouses_with_different_shipping_zone
+):
+    warehouses = warehouses_with_different_shipping_zone
+    Stock.objects.create(warehouse=warehouses[0], product_variant=variant, quantity=4)
+    Stock.objects.create(warehouse=warehouses[1], product_variant=variant, quantity=3)
+    return variant
+
+
 @pytest.fixture
 def product_variant_list(product):
     return list(
@@ -2095,6 +2126,13 @@ def warehouses_with_shipping_zone(warehouses, shipping_zone):
     return warehouses
 
 
+@pytest.fixture
+def warehouses_with_different_shipping_zone(warehouses, shipping_zones):
+    warehouses[0].shipping_zones.add(shipping_zones[0])
+    warehouses[1].shipping_zones.add(shipping_zones[1])
+    return warehouses
+
+
 @pytest.fixture
 def warehouse_no_shipping_zone(address):
     warehouse = Warehouse.objects.create(


Code review comments:
diff_hunk: @@ -527,6 +527,68 @@ def test_checkout_create_check_lines_quantity_multiple_warehouse(
     assert data["errors"][0]["field"] == "quantity"
 
 
+def test_checkout_create_sets_country_from_shipping_address_country(
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    variant_id = graphene.Node.to_global_id("ProductVariant", variant.id)
+    test_email = "test@example.com"
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = 10001
+    variables = {
+        "checkoutInput": {
+            "lines": [{"quantity": 1, "variantId": variant_id}],
+            "email": test_email,
+            "shippingAddress": shipping_address,
+        }
+    }
+    assert not Checkout.objects.exists()
+    response = user_api_client.post_graphql(MUTATION_CHECKOUT_CREATE, variables)
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutCreate"]
+    assert data["created"] is True
+    checkout = Checkout.objects.first()
+    assert checkout.country == "US"
+
+
+def test_checkout_create_check_lines_quantity_different_shipping_zone_insufficient_stocks(  # NOQA
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    stock = Stock.objects.filter(
+        warehouse__shipping_zones__countries__contains="US", product_variant=variant
+    ).first()
+    stock.quantity = 0
+    stock.save()
path: saleor/graphql/checkout/tests/test_checkout.py
text: @user1:
Maby?
```suggestion
    stock = Stock.objects.filter(
        warehouse__shipping_zones__countries__contains="US", product_variant=variant
    ).update(quantity=0)
```

diff_hunk: @@ -1278,6 +1340,92 @@ def test_checkout_shipping_address_update(
     )
 
 
+@mock.patch(
+    "saleor.graphql.checkout.mutations.update_checkout_shipping_method_if_invalid",
+    wraps=update_checkout_shipping_method_if_invalid,
+)
+def test_checkout_shipping_address_update_changes_checkout_country(
+    mocked_update_shipping_method,
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    checkout = Checkout.objects.create()
+    checkout.set_country("PL", commit=True)
+    add_variant_to_checkout(checkout, variant, 1)
+    assert checkout.shipping_address is None
+    checkout_id = graphene.Node.to_global_id("Checkout", checkout.pk)
+
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = "10001"
+    variables = {"checkoutId": checkout_id, "shippingAddress": shipping_address}
+
+    response = user_api_client.post_graphql(
+        MUTATION_CHECKOUT_SHIPPING_ADDRESS_UPDATE, variables
+    )
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutShippingAddressUpdate"]
+    assert not data["errors"]
+    checkout.refresh_from_db()
+    assert checkout.shipping_address is not None
+    assert checkout.shipping_address.first_name == shipping_address["firstName"]
+    assert checkout.shipping_address.last_name == shipping_address["lastName"]
+    assert (
+        checkout.shipping_address.street_address_1 == shipping_address["streetAddress1"]
+    )
+    assert (
+        checkout.shipping_address.street_address_2 == shipping_address["streetAddress2"]
+    )
+    assert checkout.shipping_address.postal_code == shipping_address["postalCode"]
+    assert checkout.shipping_address.country == shipping_address["country"]
+    assert checkout.shipping_address.city == shipping_address["city"].upper()
+    mocked_update_shipping_method.assert_called_once_with(
+        checkout, list(checkout), mock.ANY
+    )
+    assert checkout.country == shipping_address["country"]
+
+
+@mock.patch(
+    "saleor.graphql.checkout.mutations.update_checkout_shipping_method_if_invalid",
+    wraps=update_checkout_shipping_method_if_invalid,
+)
+def test_checkout_shipping_address_update_insufficient_stocks(
+    mocked_update_shipping_method,
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    checkout = Checkout.objects.create()
+    checkout.set_country("PL", commit=True)
+    add_variant_to_checkout(checkout, variant, 1)
+    stock = Stock.objects.filter(
+        warehouse__shipping_zones__countries__contains="US", product_variant=variant
+    ).first()
+    stock.quantity = 0
+    stock.save()
path: saleor/graphql/checkout/tests/test_checkout.py
text: @user1:
Maby?
```suggestion
    stock = Stock.objects.filter(
        warehouse__shipping_zones__countries__contains="US", product_variant=variant
    ).update(quantity=0)
```

diff_hunk: @@ -527,6 +527,68 @@ def test_checkout_create_check_lines_quantity_multiple_warehouse(
     assert data["errors"][0]["field"] == "quantity"
 
 
+def test_checkout_create_sets_country_from_shipping_address_country(
+    user_api_client,
+    variant_with_many_stocks_different_shipping_zones,
+    graphql_address_data,
+):
+    variant = variant_with_many_stocks_different_shipping_zones
+    variant_id = graphene.Node.to_global_id("ProductVariant", variant.id)
+    test_email = "test@example.com"
+    shipping_address = graphql_address_data
+    shipping_address["country"] = "US"
+    shipping_address["countryArea"] = "New York"
+    shipping_address["postalCode"] = 10001
+    variables = {
+        "checkoutInput": {
+            "lines": [{"quantity": 1, "variantId": variant_id}],
+            "email": test_email,
+            "shippingAddress": shipping_address,
+        }
+    }
+    assert not Checkout.objects.exists()
+    response = user_api_client.post_graphql(MUTATION_CHECKOUT_CREATE, variables)
+    content = get_graphql_content(response)
+    data = content["data"]["checkoutCreate"]
+    assert data["created"] is True
+    checkout = Checkout.objects.first()
+    assert checkout.country == "US"
+
+
+def test_checkout_create_check_lines_quantity_different_shipping_zone_insufficient_stocks(  # NOQA
path: saleor/graphql/checkout/tests/test_checkout.py
text: @user1:
I think we should use the test name longer than the max line length. Maybe we can use a shorter name and add some comment if still need to add more information. 

@user1:
If you set `DEFAULT_COUNTRY` in decorator before the test you can remove `different_shipping_zone` from the test name.

diff_hunk: @@ -859,6 +880,16 @@ def variant_with_many_stocks(variant, warehouses_with_shipping_zone):
     return variant
 
 
+@pytest.fixture
+def variant_with_many_stocks_different_shipping_zones(
+    variant, warehouses_with_different_shipping_zone
+):
+    warehouses = warehouses_with_different_shipping_zone
+    Stock.objects.create(warehouse=warehouses[0], product_variant=variant, quantity=4)
+    Stock.objects.create(warehouse=warehouses[1], product_variant=variant, quantity=3)
path: saleor/tests/fixtures.py
text: @user1:
Maby we can use bulk_create?