fix-code-and-test

Author: Darshan808

ipykernel/zmqshell.py

@@ -131,7 +131,9 @@ def publish(  # type:ignore[override]
                 exec_count -= 1
             outputs = getattr(self.shell.history_manager, "outputs", None)
             if outputs is not None:
-                outputs.append(HistoryOutput(output_type="display_data", bundle=data))
+                outputs.setdefault(exec_count, []).append(
+                    HistoryOutput(output_type="display_data", bundle=data)
+                )
         self._flush_streams()
         if metadata is None:
             metadata = {}

tests/test_zmq_shell.py

@@ -217,7 +217,8 @@ def test_display_stored_in_history(self):
         # Mock shell with history manager
         mock_shell = MagicMock()
         mock_shell.execution_count = 1
-        mock_shell.history_manager.outputs = {1: []}
+        mock_shell.history_manager.outputs = dict()
+        mock_shell.display_pub._in_post_execute = False
 
         self.disp_pub.shell = mock_shell