Accept a list of PIDs from deferred allowance function (#7)

Co-authored-by: Andrea Leopardi <an.leopardi@gmail.com>

lib/nimble_ownership.ex

@@ -149,6 +149,15 @@ defmodule NimbleOwnership do
   with the owner. If `pid_with_access` terminates, `pid_to_allow` will still have access to the
   key, until the `owner_pid` itself terminates or removes the allowance.
 
+  ### Deferred (lazy) allowances
+
+  If the process is not yet started at the moment of allowance definition, it might be allowed
+  as a function, assuming at the moment of invocation it would have been started.
+  If the function cannot be resolved to a PID during invocation, the expectation will not succeed.
+
+  The function might return a `t:pid/0` or a list of `t:pid/0`s. A list might be helpful
+  if one needs to allow multiple PIDs that resolve from a single term, such as the list of workers in a pool.
+
   ## Examples
 
       iex> pid = spawn(fn -> Process.sleep(:infinity) end)
@@ -161,8 +170,9 @@ defmodule NimbleOwnership do
       {:ok, self()}
 
   """
-  @spec allow(server(), pid(), pid() | (-> pid()), key()) ::
+  @spec allow(server(), pid(), pid() | (-> resolved_pid), key()) ::
           :ok | {:error, Error.t()}
+        when resolved_pid: pid() | [pid()]
   def allow(ownership_server, pid_with_access, pid_to_allow, key, timeout \\ 5000)
       when is_pid(pid_with_access) and (is_pid(pid_to_allow) or is_function(pid_to_allow, 0)) and
              is_timeout(timeout) do
@@ -610,20 +620,39 @@ defmodule NimbleOwnership do
     state.allowances
     |> Enum.reduce({[], [], false}, fn
       {key, value}, {result, resolved, unresolved} when is_function(key, 0) ->
-        case key.() do
-          pid when is_pid(pid) ->
-            {[{pid, value} | result], [{key, pid} | resolved], unresolved}
-
-          _ ->
-            {[{key, value} | result], resolved, true}
-        end
+        resolve_once(key.(), {key, value}, {result, resolved, unresolved})
 
       kv, {result, resolved, unresolved} ->
         {[kv | result], resolved, unresolved}
     end)
     |> fix_resolved(state)
   end
 
+  defp resolve_once(pid, {key, value}, {result, resolved, unresolved}) when is_pid(pid) do
+    {[{pid, value} | result], [{key, pid} | resolved], unresolved}
+  end
+
+  defp resolve_once([pid | pids], {key, value}, {result, resolved, unresolved})
+       when is_pid(pid) do
+    resolve_once(
+      pids,
+      {key, value},
+      {[{pid, value} | result], [{key, pid} | resolved], unresolved}
+    )
+  end
+
+  defp resolve_once([_not_a_pid | pids], kv, {result, resolved, _unresolved}) do
+    resolve_once(pids, kv, {[kv | result], resolved, true})
+  end
+
+  defp resolve_once([], _kv, {result, resolved, unresolved}) do
+    {result, resolved, unresolved}
+  end
+
+  defp resolve_once(_, kv, {result, resolved, _unresolved}) do
+    {[kv | result], resolved, true}
+  end
+
   defp fix_resolved({_, [], _}, state), do: state
 
   defp fix_resolved({allowances, _fun_to_pids, lazy_calls}, state) do

test/nimble_ownership_test.exs

@@ -237,6 +237,57 @@ defmodule NimbleOwnershipTest do
       assert get_meta(self(), key) == %{counter: 2}
     end
 
+    test "supports lists of lazy allowed PIDs that resolve on the next upsert", %{key: key} do
+      parent_pid = self()
+
+      # Init the key.
+      init_key(parent_pid, key, %{counter: 1})
+
+      # Allow two lazy PIDs that will resolve later.
+      assert :ok =
+               NimbleOwnership.allow(
+                 @server,
+                 self(),
+                 fn ->
+                   [Process.whereis(:lazy_pid_1), :not_a_pid, Process.whereis(:lazy_pid_2)]
+                 end,
+                 key
+               )
+
+      lazy_pid_fun =
+        fn ->
+          receive do
+            :go ->
+              assert {:ok, owner_pid} = NimbleOwnership.fetch_owner(@server, callers(), key)
+              assert owner_pid == parent_pid
+
+              NimbleOwnership.get_and_update(@server, owner_pid, key, fn info ->
+                assert %{counter: counter} = info
+                {:ok, %{counter: counter + 1}}
+              end)
+
+              send(parent_pid, :done)
+          end
+        end
+
+      {:ok, lazy_pid_1} = Task.start_link(lazy_pid_fun)
+      Process.register(lazy_pid_1, :lazy_pid_1)
+      {:ok, lazy_pid_2} = Task.start_link(lazy_pid_fun)
+      Process.register(lazy_pid_2, :lazy_pid_2)
+
+      send(lazy_pid_1, :go)
+      assert_receive :done
+
+      assert NimbleOwnership.fetch_owner(@server, [self()], key) == {:ok, self()}
+      assert get_meta(self(), key) == %{counter: 2}
+
+      send(lazy_pid_2, :go)
+      assert_receive :done
+
+      assert NimbleOwnership.fetch_owner(@server, [self()], key) == {:ok, self()}
+      assert get_meta(self(), key) == %{counter: 3}
+    end
+
     test "properly merges lazy allowed PIDs that resolve on the next upsert", %{key: key} do
       parent_pid = self()