From 5c0ad7ba1dbead715cd3501d2b7d296b33542125 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Wed, 29 Jun 2022 21:04:27 +0100
Subject: [PATCH] Fix nested frame updates (#1251)

---
 lib/livebook/notebook.ex        |  2 +-
 test/livebook/notebook_test.exs | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/lib/livebook/notebook.ex b/lib/livebook/notebook.ex
index 4668f7f8b..6d5f3e037 100644
--- a/lib/livebook/notebook.ex
+++ b/lib/livebook/notebook.ex
@@ -722,7 +722,7 @@ defmodule Livebook.Notebook do
   end
 
   defp do_find_frame_outputs({_idx, {:frame, outputs, _info}}, ref) do
-    Enum.flat_map(outputs, &find_frame_outputs(&1, ref))
+    Enum.flat_map(outputs, &do_find_frame_outputs(&1, ref))
   end
 
   defp do_find_frame_outputs(_output, _ref) do
diff --git a/test/livebook/notebook_test.exs b/test/livebook/notebook_test.exs
index 797a1d8d1..4bc7cb0a3 100644
--- a/test/livebook/notebook_test.exs
+++ b/test/livebook/notebook_test.exs
@@ -403,4 +403,29 @@ defmodule Livebook.NotebookTest do
                )
     end
   end
+
+  describe "find_frame_outputs/2" do
+    test "returns frame outputs with matching ref" do
+      frame_output = {0, {:frame, [], %{ref: "1", type: :default}}}
+
+      notebook = %{
+        Notebook.new()
+        | sections: [%{Section.new() | cells: [%{Cell.new(:code) | outputs: [frame_output]}]}]
+      }
+
+      assert [^frame_output] = Notebook.find_frame_outputs(notebook, "1")
+    end
+
+    test "finds a nested frame" do
+      nested_frame_output = {0, {:frame, [], %{ref: "2", type: :default}}}
+      frame_output = {0, {:frame, [nested_frame_output], %{ref: "1", type: :default}}}
+
+      notebook = %{
+        Notebook.new()
+        | sections: [%{Section.new() | cells: [%{Cell.new(:code) | outputs: [frame_output]}]}]
+      }
+
+      assert [^nested_frame_output] = Notebook.find_frame_outputs(notebook, "2")
+    end
+  end
 end