From 11415ce99a37ef047bb6c87736c847112395c8a2 Mon Sep 17 00:00:00 2001
From: Nav <nav@navnav.me>
Date: Thu, 19 Aug 2021 22:03:58 +0100
Subject: [PATCH] Fixed bug in TargetControllerConsole, where the
 TargetController would respond to an event before the TargetControllerConsole
 could register an interest in the response, resulting the the
 TargetControllerConsole missing the response.

---
 src/TargetController/TargetController.hpp     |  2 +-
 .../TargetControllerConsole.cpp               | 25 +++-----
 .../TargetControllerConsole.hpp               | 57 +++++++++++++++++++
 3 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/src/TargetController/TargetController.hpp b/src/TargetController/TargetController.hpp
index efa00b74..6f4edd78 100644
--- a/src/TargetController/TargetController.hpp
+++ b/src/TargetController/TargetController.hpp
@@ -104,7 +104,7 @@ namespace Bloom
          * @return
          */
         static auto getSupportedTargets() {
-            static std::map<std::string, std::function<std::unique_ptr<Targets::Target>()>> mapping;
+            static auto mapping = std::map<std::string, std::function<std::unique_ptr<Targets::Target>()>>();
 
             if (mapping.empty()) {
                 mapping = {
diff --git a/src/TargetController/TargetControllerConsole.cpp b/src/TargetController/TargetControllerConsole.cpp
index 9460063b..2145a236 100644
--- a/src/TargetController/TargetControllerConsole.cpp
+++ b/src/TargetController/TargetControllerConsole.cpp
@@ -1,6 +1,8 @@
-#include <cstdint>
-
 #include "TargetControllerConsole.hpp"
+
+#include <cstdint>
+#include <thread>
+
 #include "src/EventManager/Events/Events.hpp"
 #include "src/Logger/Logger.hpp"
 
@@ -37,21 +39,10 @@ bool TargetControllerConsole::isTargetControllerInService() noexcept {
 }
 
 Targets::TargetDescriptor TargetControllerConsole::getTargetDescriptor() {
-    auto extractEvent = std::make_shared<ExtractTargetDescriptor>();
-    this->eventManager.triggerEvent(extractEvent);
-    auto responseEvent = this->eventListener.waitForEvent<
-        TargetDescriptorExtracted,
-        TargetControllerErrorOccurred
-    >(this->defaultTimeout, extractEvent->id);
-
-    if (!responseEvent.has_value()
-        || !std::holds_alternative<SharedEventPointer<TargetDescriptorExtracted>>(responseEvent.value())
-    ) {
-        throw Exception("Unexpected response from TargetController");
-    }
-
-    auto descriptorExtracted = std::get<SharedEventPointer<TargetDescriptorExtracted>>(responseEvent.value());
-    return descriptorExtracted->targetDescriptor;
+    auto descriptorExtractedEvent = this->triggerTargetControllerEventAndWaitForResponse(
+        std::make_shared<ExtractTargetDescriptor>()
+    );
+    return descriptorExtractedEvent->targetDescriptor;
 }
 
 void TargetControllerConsole::stopTargetExecution() {
diff --git a/src/TargetController/TargetControllerConsole.hpp b/src/TargetController/TargetControllerConsole.hpp
index 34aeeda1..87ad0db9 100644
--- a/src/TargetController/TargetControllerConsole.hpp
+++ b/src/TargetController/TargetControllerConsole.hpp
@@ -29,6 +29,63 @@ namespace Bloom
 
         std::chrono::milliseconds defaultTimeout = std::chrono::milliseconds(10000);
 
+        /**
+         * Triggers an event for the TargetController and waits for a response.
+         *
+         * To use this method, the triggered event must define a 'TargetControllerResponseType' alias, which should
+         * specify the type of response expected by the TargetController.
+         * For an example of this, see the Events::ExtractTargetDescriptor class.
+         *
+         * If the TargetController fails to respond within the given time specified by the timeout parameter, or it
+         * responds with an instance of Events::TargetControllerErrorOccurred, this method will throw an exception.
+         *
+         * @tparam TriggerEventType
+         *
+         * @param event
+         *  Event to trigger.
+         *
+         * @param timeout
+         *  The time, in milliseconds, to wait for the TargetController to respond to the event. If this is not
+         *  supplied, this->defaultTimeout will be used.
+         *
+         * @return
+         */
+        template<class TriggerEventType>
+        auto triggerTargetControllerEventAndWaitForResponse(
+            const Events::SharedEventPointerNonConst<TriggerEventType> event,
+            std::optional<std::chrono::milliseconds> timeout = {}
+        ) {
+            using ResponseEventType = typename TriggerEventType::TargetControllerResponseType;
+
+            bool deRegisterEventType = false;
+
+            if (!this->eventListener.isEventTypeRegistered<ResponseEventType>()) {
+                this->eventListener.registerEventType<ResponseEventType>();
+                deRegisterEventType = true;
+            }
+
+            this->eventManager.triggerEvent(event);
+
+            auto responseEvent = this->eventListener.waitForEvent<
+                ResponseEventType,
+                Events::TargetControllerErrorOccurred
+            >(timeout.value_or(this->defaultTimeout), event->id);
+
+            if (deRegisterEventType) {
+                this->eventListener.deRegisterEventType<ResponseEventType>();
+            }
+
+            if (!responseEvent.has_value()) {
+                throw Bloom::Exceptions::Exception("Timed out waiting for response from TargetController.");
+            }
+
+            if (!std::holds_alternative<Events::SharedEventPointer<ResponseEventType>>(responseEvent.value())) {
+                throw Bloom::Exceptions::Exception("Unexpected response from TargetController");
+            }
+
+            return std::get<Events::SharedEventPointer<ResponseEventType>>(responseEvent.value());
+        }
+
     public:
         TargetControllerConsole(EventManager& eventManager, EventListener& eventListener):
         eventManager(eventManager), eventListener(eventListener) {};