Ecto cast_assoc partial set

_test.exs

defmodule TestEctoReplace.ContextTest do
  use TestEctoReplace.DataCase
  alias TestEctoReplace.Context.{Country, Restaurant}

  test "assoc" do
    ##################
    # Arrange
    country = Repo.insert!(%Country{name: "Germany"})

    [_, _, kfc, burger_king, _] =
      ["Mc Donalds", "Vapiano", "KFC", "Burger King", "Wienerwald"]
      |> Enum.map(fn name ->
        %Restaurant{name: name, country_id: country.id}
        |> Repo.insert!()
      end)

    # Page 2 for a limit: 2
    paginated_ids = [kfc.id, burger_king.id]

    ##################
    # Act
    query = from(r in Restaurant, where: r.id in ^paginated_ids)

    Repo.get(Country, country.id)
    |> Repo.preload(restaurants: query)
    |> Country.changeset(%{
      restaurants: [
        %{id: kfc.id, name: "KFC Changed"},
        %{id: burger_king.id, name: "Burger King Changed"}
      ]
    })
    |> Repo.update!()

    ##################
    # Assert
    restaurants =
      Repo.get(Country, country.id)
      |> Repo.preload(:restaurants)
      |> Map.get(:restaurants)

    assert [
             %{name: "Mc Donalds"},
             %{name: "Vapiano"},
             %{name: "Wienerwald"},
             %{name: "KFC Changed"},
             %{name: "Burger King Changed"}
           ] = restaurants
  end
end

migration.exs

defmodule TestEctoReplace.Repo.Migrations.CreateTables do
  use Ecto.Migration

  def change do
    create table(:countries) do
      add(:name, :string, null: false)
      timestamps()
    end

    create table(:restaurants) do
      add(:name, :string, null: false)
      add(:country_id, references(:countries), null: false)
      timestamps()
    end
  end
end

schemas.ex

defmodule TestEctoReplace.Context.Restaurant do
  use Ecto.Schema
  import Ecto.Changeset

  schema "restaurants" do
    field(:name, :string)

    belongs_to(:country, TestEctoReplace.Context.Country)

    timestamps()
  end

  def changeset(struct, params) do
    struct
    |> cast(params, [:name])
    |> validate_required([:name])
  end
end

defmodule TestEctoReplace.Context.Country do
  use Ecto.Schema
  import Ecto.Changeset

  schema "countries" do
    field(:name, :string)

    has_many(:restaurants, TestEctoReplace.Context.Restaurant)

    timestamps()
  end

  def changeset(struct, params) do
    struct
    |> cast(params, [:name])
    |> validate_required([:name])
    |> cast_assoc(:restaurants)
  end
end