From d6b379fddd6098a8ec64c9566c762631187e570e Mon Sep 17 00:00:00 2001
From: Aleksei Matiushkin <aleksei.matiushkin@kantox.com>
Date: Mon, 22 Jul 2024 19:03:47 +0200
Subject: [PATCH] Fix several allowances with same keys (#6)

---
 lib/nimble_ownership.ex        |  7 +++-
 test/nimble_ownership_test.exs | 64 ++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/lib/nimble_ownership.ex b/lib/nimble_ownership.ex
index c1999da..12b8fad 100644
--- a/lib/nimble_ownership.ex
+++ b/lib/nimble_ownership.ex
@@ -627,6 +627,11 @@ defmodule NimbleOwnership do
   defp fix_resolved({_, [], _}, state), do: state
 
   defp fix_resolved({allowances, _fun_to_pids, lazy_calls}, state) do
-    %__MODULE__{state | allowances: Map.new(allowances), lazy_calls: lazy_calls}
+    allowances =
+      Enum.reduce(allowances, %{}, fn {k, v}, acc ->
+        Map.update(acc, k, v, &Map.merge(&1, v))
+      end)
+
+    %__MODULE__{state | allowances: allowances, lazy_calls: lazy_calls}
   end
 end
diff --git a/test/nimble_ownership_test.exs b/test/nimble_ownership_test.exs
index eacfa28..a803fd2 100644
--- a/test/nimble_ownership_test.exs
+++ b/test/nimble_ownership_test.exs
@@ -237,6 +237,70 @@ defmodule NimbleOwnershipTest do
       assert get_meta(self(), key) == %{counter: 2}
     end
 
+    test "properly merges lazy allowed PIDs that resolve on the next upsert", %{key: key} do
+      parent_pid = self()
+
+      parent_process_fun = fn counter ->
+        fn ->
+          # Needed to trick double allowance checker 
+          Process.delete(:"$callers")
+
+          # Init the key
+          key = "#{counter} → #{key}"
+          init_key(self(), key, %{counter: 1})
+
+          # Allow two lazy PID that will resolve later
+          assert :ok =
+                   NimbleOwnership.allow(
+                     @server,
+                     self(),
+                     fn -> Process.whereis(:lazy_pid) end,
+                     key
+                   )
+
+          receive do
+            {:go, lazy_pid} ->
+              send(lazy_pid, {:go, self(), key})
+              assert_receive :done
+
+              assert NimbleOwnership.fetch_owner(@server, [self()], key) == {:ok, self()}
+              assert get_meta(self(), key) == %{counter: 2}
+
+              send(parent_pid, :parent_process_done)
+          end
+        end
+      end
+
+      {:ok, pid_1} = Task.start_link(parent_process_fun.(1))
+      {:ok, pid_2} = Task.start_link(parent_process_fun.(2))
+
+      lazy_process_fun = fn ->
+        for _ <- 1..2 do
+          receive do
+            {:go, parent_pid, key} ->
+              assert {:ok, owner_pid} = NimbleOwnership.fetch_owner(@server, [parent_pid], key)
+              assert owner_pid == parent_pid
+
+              NimbleOwnership.get_and_update(@server, owner_pid, key, fn info ->
+                assert info == %{counter: 1}
+                {:ok, %{counter: 2}}
+              end)
+
+              send(parent_pid, :done)
+          end
+        end
+      end
+
+      {:ok, lazy_pid} = Task.start_link(lazy_process_fun)
+      Process.register(lazy_pid, :lazy_pid)
+
+      send(pid_1, {:go, lazy_pid})
+      send(pid_2, {:go, lazy_pid})
+
+      assert_receive :parent_process_done
+      assert_receive :parent_process_done
+    end
+
     test "ignores lazy PIDs that don't actually resolve to a PID", %{key: key} do
       owner_pid = self()