From c204aa1abe9818a2624c4cb62c20616c20bb2eb9 Mon Sep 17 00:00:00 2001
From: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>
Date: Tue, 28 Sep 2021 18:16:02 -0700
Subject: [PATCH] Change notification buses to be addressed on EntityContextId,
 and connect the handlers to the default id. Request interfaces should not
 need that id as they can retrieve it from the entityId.

Signed-off-by: Danilo Aimini <82231674+AMZN-daimini@users.noreply.github.com>
---
 .../FocusMode/FocusModeNotificationBus.h          | 15 ++++++++++++---
 .../Prefab/PrefabFocusHandler.cpp                 |  4 ++++
 .../Prefab/PrefabFocusNotificationBus.h           | 13 ++++++++++---
 .../UI/Outliner/EntityOutlinerTreeView.cpp        |  7 ++++++-
 .../UI/Prefab/PrefabViewportFocusPathHandler.cpp  | 10 ++++++++--
 5 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/FocusMode/FocusModeNotificationBus.h b/Code/Framework/AzToolsFramework/AzToolsFramework/FocusMode/FocusModeNotificationBus.h
index a16bb9a3ef..ab8629ac85 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/FocusMode/FocusModeNotificationBus.h
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/FocusMode/FocusModeNotificationBus.h
@@ -11,19 +11,28 @@
 #include <AzCore/Component/EntityId.h>
 #include <AzCore/EBus/EBus.h>
 
+#include <AzFramework/Entity/EntityContext.h>
+
 namespace AzToolsFramework
 {
     //! Used to notify when the editor focus changes.
     class FocusModeNotifications
         : public AZ::EBusTraits
     {
-    protected:
-        ~FocusModeNotifications() = default;
-
     public:
+        //////////////////////////////////////////////////////////////////////////
+        // EBusTraits overrides
+        static const AZ::EBusHandlerPolicy HandlerPolicy = AZ::EBusHandlerPolicy::Multiple;
+        static const AZ::EBusAddressPolicy AddressPolicy = AZ::EBusAddressPolicy::ById;
+        using BusIdType = AzFramework::EntityContextId;
+        //////////////////////////////////////////////////////////////////////////
+        
         //! Triggered when the editor focus is changed to a different entity.
         //! @param entityId The entity the focus has been moved to.
         virtual void OnEditorFocusChanged(AZ::EntityId entityId) = 0;
+
+    protected:
+        ~FocusModeNotifications() = default;
     };
 
     using FocusModeNotificationBus = AZ::EBus<FocusModeNotifications>;
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusHandler.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusHandler.cpp
index 9bfe64bc84..3ce8065ae5 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusHandler.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusHandler.cpp
@@ -96,6 +96,10 @@ namespace AzToolsFramework::Prefab
             else
             {
                 containerEntityId = AZ::EntityId();
+
+                // Clear the selection
+                AzToolsFramework::SelectEntities({});
+                
             }
             
             // Focus on the descendants of the container entity
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusNotificationBus.h b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusNotificationBus.h
index aa27d87013..31bab719e2 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusNotificationBus.h
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabFocusNotificationBus.h
@@ -17,12 +17,19 @@ namespace AzToolsFramework::Prefab
     class PrefabFocusNotifications
         : public AZ::EBusTraits
     {
-    protected:
-        ~PrefabFocusNotifications() = default;
-
     public:
+        //////////////////////////////////////////////////////////////////////////
+        // EBusTraits overrides
+        static const AZ::EBusHandlerPolicy HandlerPolicy = AZ::EBusHandlerPolicy::Multiple;
+        static const AZ::EBusAddressPolicy AddressPolicy = AZ::EBusAddressPolicy::ById;
+        using BusIdType = AzFramework::EntityContextId;
+        //////////////////////////////////////////////////////////////////////////
+        
         //! Triggered when the editor focus is changed to a different prefab.
         virtual void OnPrefabFocusChanged() = 0;
+
+    protected:
+        ~PrefabFocusNotifications() = default;
     };
 
     using PrefabFocusNotificationBus = AZ::EBus<PrefabFocusNotifications>;
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Outliner/EntityOutlinerTreeView.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Outliner/EntityOutlinerTreeView.cpp
index 6a40b3be3a..4364ae1efc 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Outliner/EntityOutlinerTreeView.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Outliner/EntityOutlinerTreeView.cpp
@@ -39,7 +39,12 @@ namespace AzToolsFramework
         AZ_Assert((m_editorEntityFrameworkInterface != nullptr),
             "EntityOutlinerTreeView requires a EditorEntityFrameworkInterface instance on Construction.");
 
-        FocusModeNotificationBus::Handler::BusConnect();
+        
+        AzFramework::EntityContextId editorEntityContextId = AzFramework::EntityContextId::CreateNull();
+        AzToolsFramework::EditorEntityContextRequestBus::BroadcastResult(
+            editorEntityContextId, &AzToolsFramework::EditorEntityContextRequestBus::Events::GetEditorEntityContextId);
+
+        FocusModeNotificationBus::Handler::BusConnect(editorEntityContextId);
 
         viewport()->setMouseTracking(true);
     }
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabViewportFocusPathHandler.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabViewportFocusPathHandler.cpp
index c1dd66ec7f..09d0579dab 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabViewportFocusPathHandler.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabViewportFocusPathHandler.cpp
@@ -6,15 +6,21 @@
  *
  */
 
-#include <AzToolsFramework/Prefab/PrefabFocusInterface.h>
 #include <AzToolsFramework/UI/Prefab/PrefabViewportFocusPathHandler.h>
 
+#include <AzToolsFramework/Entity/EditorEntityContextBus.h>
+#include <AzToolsFramework/Prefab/PrefabFocusInterface.h>
+
 namespace AzToolsFramework::Prefab
 {
     PrefabViewportFocusPathHandler::PrefabViewportFocusPathHandler()
     {
         // Connect to Prefab Focus Notifications
-        PrefabFocusNotificationBus::Handler::BusConnect();
+        AzFramework::EntityContextId editorEntityContextId = AzFramework::EntityContextId::CreateNull();
+        AzToolsFramework::EditorEntityContextRequestBus::BroadcastResult(
+            editorEntityContextId, &AzToolsFramework::EditorEntityContextRequestBus::Events::GetEditorEntityContextId);
+
+        PrefabFocusNotificationBus::Handler::BusConnect(editorEntityContextId);
     }
 
     PrefabViewportFocusPathHandler::~PrefabViewportFocusPathHandler()