From f27946956423a646359c2f70a303c646711dc5fe Mon Sep 17 00:00:00 2001
From: Anthony Drendel <me@anthonydrendel.com>
Date: Wed, 20 Nov 2024 02:26:33 +0100
Subject: [PATCH] Add ExMachina.EctoPolymorphicEmbed

---
 lib/ex_machina/ecto_polymorphic_embed.ex      | 380 ++++++++++++++++++
 .../ecto_polymorphic_embed_strategy.ex        | 155 +++++++
 2 files changed, 535 insertions(+)
 create mode 100644 lib/ex_machina/ecto_polymorphic_embed.ex
 create mode 100644 lib/ex_machina/ecto_polymorphic_embed_strategy.ex

diff --git a/lib/ex_machina/ecto_polymorphic_embed.ex b/lib/ex_machina/ecto_polymorphic_embed.ex
new file mode 100644
index 0000000..a12cf6f
--- /dev/null
+++ b/lib/ex_machina/ecto_polymorphic_embed.ex
@@ -0,0 +1,380 @@
+defmodule ExMachina.EctoPolymorphicEmbed do
+  @moduledoc """
+  Module for building and inserting factories with Ecto
+
+  This module works much like the regular `ExMachina` module, but adds a few
+  nice things that make working with Ecto easier.
+
+  * It uses `ExMachina.EctoPolymorphicEmbedStrategy`, which adds `insert/1`, `insert/2`,
+  `insert/3` `insert_pair/2`, `insert_list/3`.
+  * Adds a `params_for` function that is useful for working with changesets or
+    sending params to API endpoints.
+
+  More in-depth examples are in the [README](readme.html).
+  """
+
+  @callback insert(factory_name :: atom) :: any
+  @callback insert(factory_name :: atom, attrs :: keyword | map) :: any
+
+  @doc """
+  Builds a factory and inserts it into the database.
+
+  The first two arguments are the same as `c:ExMachina.build/2`. The last
+  argument is a set of options that will be passed to Ecto's
+  [`Repo.insert!/2`](https://hexdocs.pm/ecto/Ecto.Repo.html#c:insert!/2).
+
+  ## Examples
+
+      # return all values from the database
+      insert(:user, [name: "Jane"], returning: true)
+      build(:user, name: "Jane") |> insert(returning: true)
+
+      # use a different prefix
+      insert(:user, [name: "Jane"], prefix: "other_tenant")
+      build(:user, name: "Jane") |> insert(prefix: "other_tenant")
+  """
+  @callback insert(factory_name :: atom, attrs :: keyword | map, opts :: keyword | map) :: any
+
+  @doc """
+  Builds two factories and inserts them into the database.
+
+  The arguments are the same as `c:ExMachina.build_pair/2`.
+  """
+  @callback insert_pair(factory_name :: atom) :: list
+  @callback insert_pair(factory_name :: atom, attrs :: keyword | map) :: list
+
+  @doc """
+  Builds many factories and inserts them into the database.
+
+  The arguments are the same as `c:ExMachina.build_list/3`.
+  """
+  @callback insert_list(number_of_records :: integer, factory_name :: atom) :: list
+  @callback insert_list(
+              number_of_records :: integer,
+              factory_name :: atom,
+              attrs :: keyword | map
+            ) :: list
+
+  @doc """
+  Builds a factory and returns only its fields.
+
+  This is only for use with Ecto models.
+
+  Will return a map with the fields and virtual fields, but without the Ecto
+  metadata, the primary key, or any `belongs_to` associations. This will
+  recursively act on `has_one` associations and Ecto structs found in
+  `has_many` associations.
+
+  If you want `belongs_to` associations to be inserted, use
+  `c:params_with_assocs/2`.
+
+  If you want params with string keys use `c:string_params_for/2`.
+
+  ## Example
+
+      def user_factory do
+        %MyApp.User{name: "John Doe", admin: false}
+      end
+
+      # Returns %{name: "John Doe", admin: true}
+      params_for(:user, admin: true)
+
+      # Returns %{name: "John Doe", admin: false}
+      params_for(:user)
+  """
+  @callback params_for(factory_name :: atom) :: %{optional(atom) => any}
+  @callback params_for(factory_name :: atom, attrs :: keyword | map) :: %{optional(atom) => any}
+
+  @doc """
+  Similar to `c:params_for/2` but converts atom keys to strings in returned map.
+
+  The result of this function can be safely used in controller tests for Phoenix
+  web applications.
+
+  ## Example
+
+      def user_factory do
+        %MyApp.User{name: "John Doe", admin: false}
+      end
+
+      # Returns %{"name" => "John Doe", "admin" => true}
+      string_params_for(:user, admin: true)
+  """
+  @callback string_params_for(factory_name :: atom) :: %{optional(String.t()) => any}
+  @callback string_params_for(factory_name :: atom, attrs :: keyword | map) :: %{
+              optional(String.t()) => any
+            }
+
+  @doc """
+  Similar to `c:params_for/2` but inserts all `belongs_to` associations and
+  sets the foreign keys.
+
+  If you want params with string keys use `c:string_params_with_assocs/2`.
+
+  ## Example
+
+      def article_factory do
+        %MyApp.Article{title: "An Awesome Article", author: build(:author)}
+      end
+
+      # Inserts an author and returns %{title: "An Awesome Article", author_id: 12}
+      params_with_assocs(:article)
+  """
+  @callback params_with_assocs(factory_name :: atom) :: %{optional(atom) => any}
+  @callback params_with_assocs(factory_name :: atom, attrs :: keyword | map) :: %{
+              optional(atom) => any
+            }
+  @doc """
+  Similar to `c:params_with_assocs/2` but converts atom keys to strings in
+  returned map.
+
+  The result of this function can be safely used in controller tests for Phoenix
+  web applications.
+
+  ## Example
+
+      def article_factory do
+        %MyApp.Article{title: "An Awesome Article", author: build(:author)}
+      end
+
+      # Inserts an author and returns %{"title" => "An Awesome Article", "author_id" => 12}
+      string_params_with_assocs(:article)
+  """
+  @callback string_params_with_assocs(factory_name :: atom) :: %{optional(String.t()) => any}
+  @callback string_params_with_assocs(factory_name :: atom, attrs :: keyword | map) :: %{
+              optional(String.t()) => any
+            }
+
+  defmacro __using__(opts) do
+    verify_ecto_dep()
+
+    quote do
+      use ExMachina
+      use ExMachina.EctoPolymorphicEmbedStrategy, repo: unquote(Keyword.get(opts, :repo))
+
+      def params_for(factory_name, attrs \\ %{}) do
+        ExMachina.EctoPolymorphicEmbed.params_for(__MODULE__, factory_name, attrs)
+      end
+
+      def string_params_for(factory_name, attrs \\ %{}) do
+        ExMachina.EctoPolymorphicEmbed.string_params_for(__MODULE__, factory_name, attrs)
+      end
+
+      def params_with_assocs(factory_name, attrs \\ %{}) do
+        ExMachina.EctoPolymorphicEmbed.params_with_assocs(__MODULE__, factory_name, attrs)
+      end
+
+      def string_params_with_assocs(factory_name, attrs \\ %{}) do
+        ExMachina.EctoPolymorphicEmbed.string_params_with_assocs(__MODULE__, factory_name, attrs)
+      end
+    end
+  end
+
+  @doc false
+  def params_for(module, factory_name, attrs \\ %{}) do
+    factory_name
+    |> module.build(attrs)
+    |> recursively_strip
+  end
+
+  @doc false
+  def string_params_for(module, factory_name, attrs \\ %{}) do
+    module
+    |> params_for(factory_name, attrs)
+    |> convert_atom_keys_to_strings
+  end
+
+  @doc false
+  def params_with_assocs(module, factory_name, attrs \\ %{}) do
+    factory_name
+    |> module.build(attrs)
+    |> insert_belongs_to_assocs(module)
+    |> recursively_strip
+  end
+
+  @doc false
+  def string_params_with_assocs(module, factory_name, attrs \\ %{}) do
+    module
+    |> params_with_assocs(factory_name, attrs)
+    |> convert_atom_keys_to_strings
+  end
+
+  defp recursively_strip(%{__struct__: _} = record) do
+    record
+    |> set_persisted_belongs_to_ids
+    |> handle_assocs
+    |> handle_embeds
+    |> drop_ecto_fields
+    |> drop_fields_with_nil_values
+  end
+
+  defp recursively_strip(record), do: record
+
+  defp handle_assocs(%{__struct__: struct} = record) do
+    associations = struct.__schema__(:associations)
+
+    Enum.reduce(associations, record, fn association_name, record ->
+      case struct.__schema__(:association, association_name) do
+        %{__struct__: Ecto.Association.BelongsTo} ->
+          Map.delete(record, association_name)
+
+        _ ->
+          record
+          |> Map.get(association_name)
+          |> handle_assoc(record, association_name)
+      end
+    end)
+  end
+
+  defp handle_assoc(original_assoc, record, association_name) do
+    case original_assoc do
+      %{__meta__: %{__struct__: Ecto.Schema.Metadata, state: :built}} ->
+        assoc = recursively_strip(original_assoc)
+        Map.put(record, association_name, assoc)
+
+      nil ->
+        Map.put(record, association_name, nil)
+
+      list when is_list(list) ->
+        has_many_assoc = Enum.map(original_assoc, &recursively_strip/1)
+        Map.put(record, association_name, has_many_assoc)
+
+      %{__struct__: Ecto.Association.NotLoaded} ->
+        Map.delete(record, association_name)
+    end
+  end
+
+  defp handle_embeds(%{__struct__: struct} = record) do
+    embeds = struct.__schema__(:embeds)
+
+    Enum.reduce(embeds, record, fn embed_name, record ->
+      record
+      |> Map.get(embed_name)
+      |> handle_embed(record, embed_name)
+    end)
+  end
+
+  defp handle_embed(original_embed, record, embed_name) do
+    case original_embed do
+      %{} ->
+        embed = recursively_strip(original_embed)
+        Map.put(record, embed_name, embed)
+
+      list when is_list(list) ->
+        embeds_many = Enum.map(original_embed, &recursively_strip/1)
+        Map.put(record, embed_name, embeds_many)
+
+      nil ->
+        Map.delete(record, embed_name)
+    end
+  end
+
+  defp set_persisted_belongs_to_ids(%{__struct__: struct} = record) do
+    associations = struct.__schema__(:associations)
+
+    Enum.reduce(associations, record, fn association_name, record ->
+      association = struct.__schema__(:association, association_name)
+
+      with %{__struct__: Ecto.Association.BelongsTo} <- association,
+           belongs_to <- Map.get(record, association_name),
+           %{__meta__: %{__struct__: Ecto.Schema.Metadata, state: :loaded}} <- belongs_to do
+        set_belongs_to_primary_key(record, belongs_to, association)
+      else
+        _ -> record
+      end
+    end)
+  end
+
+  defp set_belongs_to_primary_key(record, belongs_to, association) do
+    primary_key = Map.get(belongs_to, association.related_key)
+    Map.put(record, association.owner_key, primary_key)
+  end
+
+  defp insert_belongs_to_assocs(%{__struct__: struct} = record, module) do
+    associations = struct.__schema__(:associations)
+
+    Enum.reduce(associations, record, fn association_name, record ->
+      case struct.__schema__(:association, association_name) do
+        association = %{__struct__: Ecto.Association.BelongsTo} ->
+          insert_built_belongs_to_assoc(module, association, record)
+
+        _ ->
+          record
+      end
+    end)
+  end
+
+  defp insert_built_belongs_to_assoc(module, association, record) do
+    case Map.get(record, association.field) do
+      built_relation = %{__meta__: %{state: :built}} ->
+        relation = module.insert(built_relation)
+        set_belongs_to_primary_key(record, relation, association)
+
+      _ ->
+        Map.delete(record, association.owner_key)
+    end
+  end
+
+  @doc false
+  def drop_ecto_fields(%{__struct__: struct} = record) do
+    record
+    |> Map.from_struct()
+    |> Map.delete(:__meta__)
+    |> drop_autogenerated_ids(struct)
+  end
+
+  def drop_ecto_fields(embedded_record), do: embedded_record
+
+  defp drop_autogenerated_ids(map, struct) do
+    case struct.__schema__(:autogenerate_id) do
+      {name, _source, _type} -> Map.delete(map, name)
+      {name, _type} -> Map.delete(map, name)
+      nil -> map
+    end
+  end
+
+  defp drop_fields_with_nil_values(map) do
+    map
+    |> Enum.reject(fn {_, value} -> value == nil end)
+    |> Enum.into(%{})
+  end
+
+  defp convert_atom_keys_to_strings(values) when is_list(values) do
+    Enum.map(values, &convert_atom_keys_to_strings/1)
+  end
+
+  defp convert_atom_keys_to_strings(%NaiveDateTime{} = value) do
+    if Application.get_env(:ex_machina, :preserve_dates, false) do
+      value
+    else
+      value |> Map.from_struct() |> convert_atom_keys_to_strings()
+    end
+  end
+
+  defp convert_atom_keys_to_strings(%DateTime{} = value) do
+    if Application.get_env(:ex_machina, :preserve_dates, false) do
+      value
+    else
+      value |> Map.from_struct() |> convert_atom_keys_to_strings()
+    end
+  end
+
+  defp convert_atom_keys_to_strings(%{__struct__: _} = record) when is_map(record) do
+    record |> Map.from_struct() |> convert_atom_keys_to_strings()
+  end
+
+  defp convert_atom_keys_to_strings(record) when is_map(record) do
+    Enum.reduce(record, Map.new(), fn {key, value}, acc ->
+      Map.put(acc, to_string(key), convert_atom_keys_to_strings(value))
+    end)
+  end
+
+  defp convert_atom_keys_to_strings(value), do: value
+
+  defp verify_ecto_dep do
+    unless Code.ensure_loaded?(Ecto) do
+      raise "You tried to use ExMachina.EctoPolymorphicEmbed, but the Ecto module is not loaded. " <>
+              "Please add ecto to your dependencies."
+    end
+  end
+end
diff --git a/lib/ex_machina/ecto_polymorphic_embed_strategy.ex b/lib/ex_machina/ecto_polymorphic_embed_strategy.ex
new file mode 100644
index 0000000..1e4d0e4
--- /dev/null
+++ b/lib/ex_machina/ecto_polymorphic_embed_strategy.ex
@@ -0,0 +1,155 @@
+defmodule ExMachina.EctoPolymorphicEmbedStrategy do
+  @moduledoc false
+
+  # Copied from https://github.com/thoughtbot/ex_machina/blob/b2f47a36b84fded6c37434bbf9041b33b30387e9/lib/ex_machina/ecto_strategy.ex
+
+  use ExMachina.Strategy, function_name: :insert
+
+  def handle_insert(%{__meta__: %{__struct__: Ecto.Schema.Metadata}} = record, %{repo: repo}) do
+    record
+    |> cast
+    |> repo.insert!
+  end
+
+  def handle_insert(resource, options),
+    do: ExMachina.EctoStrategy.handle_insert(resource, options)
+
+  defp cast(record) do
+    record
+    |> cast_all_fields
+    |> cast_all_embeds
+    ## Change Start
+    |> cast_polymorphic_embeds
+    ## Change End
+    |> cast_all_assocs
+  end
+
+  defp cast_all_fields(%{__struct__: schema} = struct) do
+    schema
+    |> schema_fields()
+    |> Enum.reduce(struct, fn field_key, struct ->
+      casted_value = cast_field(field_key, struct)
+
+      Map.put(struct, field_key, casted_value)
+    end)
+  end
+
+  defp cast_field(field_key, %{__struct__: schema} = struct) do
+    field_type = schema.__schema__(:type, field_key)
+    value = Map.get(struct, field_key)
+
+    cast_value(field_type, value, struct)
+  end
+
+  defp cast_value(field_type, value, struct) do
+    case Ecto.Type.cast(field_type, value) do
+      {:ok, value} ->
+        value
+
+      _ ->
+        raise "Failed to cast `#{inspect(value)}` of type #{inspect(field_type)} in #{
+                inspect(struct)
+              }."
+    end
+  end
+
+  defp cast_all_embeds(%{__struct__: schema} = struct) do
+    schema
+    |> schema_embeds()
+    |> Enum.reduce(struct, fn embed_key, struct ->
+      casted_value = struct |> Map.get(embed_key) |> cast_embed(embed_key, struct)
+
+      Map.put(struct, embed_key, casted_value)
+    end)
+  end
+
+  defp cast_embed(embeds_many, embed_key, struct) when is_list(embeds_many) do
+    Enum.map(embeds_many, &cast_embed(&1, embed_key, struct))
+  end
+
+  defp cast_embed(embed, embed_key, %{__struct__: schema}) do
+    if embed do
+      embedding_reflection = schema.__schema__(:embed, embed_key)
+      embed_type = embedding_reflection.related
+      embed_type |> struct() |> Map.merge(embed) |> cast()
+    end
+  end
+
+  ## Change Start
+  defp cast_polymorphic_embeds(%{__struct__: schema} = struct) do
+    schema
+    |> schema_polymorphic_embeds()
+    |> Enum.reduce(struct, fn embed_key, struct ->
+      casted_value = struct |> Map.get(embed_key)
+
+      Map.put(struct, embed_key, casted_value)
+    end)
+  end
+
+  ## Change End
+
+  defp cast_all_assocs(%{__struct__: schema} = struct) do
+    assoc_keys = schema_associations(schema)
+
+    Enum.reduce(assoc_keys, struct, fn assoc_key, struct ->
+      casted_value = struct |> Map.get(assoc_key) |> cast_assoc(assoc_key, struct)
+
+      Map.put(struct, assoc_key, casted_value)
+    end)
+  end
+
+  defp cast_assoc(has_many_assoc, assoc_key, struct) when is_list(has_many_assoc) do
+    Enum.map(has_many_assoc, &cast_assoc(&1, assoc_key, struct))
+  end
+
+  defp cast_assoc(assoc, assoc_key, %{__struct__: schema}) do
+    case assoc do
+      %{__meta__: %{__struct__: Ecto.Schema.Metadata, state: :built}} ->
+        cast(assoc)
+
+      %{__struct__: Ecto.Association.NotLoaded} ->
+        assoc
+
+      %{__struct__: _} ->
+        cast(assoc)
+
+      %{} ->
+        assoc_reflection = schema.__schema__(:association, assoc_key)
+        assoc_type = assoc_reflection.related
+        assoc_type |> struct() |> Map.merge(assoc) |> cast()
+
+      nil ->
+        nil
+    end
+  end
+
+  defp schema_fields(schema) do
+    ## Change Start
+    (schema_non_virtual_fields(schema) -- schema_embeds(schema)) --
+      schema_polymorphic_embeds(schema)
+
+    ## Change End
+  end
+
+  defp schema_non_virtual_fields(schema) do
+    schema.__schema__(:fields)
+  end
+
+  ## Change Start
+  defp schema_polymorphic_embeds(schema) do
+    Enum.filter(
+      schema.__schema__(:fields),
+      &match?({:parameterized, PolymorphicEmbed, _options}, schema.__schema__(:type, &1))
+    )
+  end
+
+  ## Change End
+
+  defp schema_embeds(schema) do
+    schema.__schema__(:embeds)
+  end
+
+  defp schema_associations(schema) do
+    schema.__schema__(:associations)
+  end
+end
\ No newline at end of file