From fc3747061c0d78b345c068da72e8366f7dd8a7c9 Mon Sep 17 00:00:00 2001
From: Henri Casanova <henricasanova@gmail.com>
Date: Thu, 16 Jul 2020 09:35:30 -1000
Subject: [PATCH] (#187) Added a CloudComputeService::getVMComputeService
 method

---
 .../compute/cloud/CloudComputeService.h       | 10 ++++++----
 .../compute/cloud/CloudComputeService.cpp     | 19 +++++++++++++++++++
 .../VirtualizedClusterServiceTest.cpp         | 11 +++++++++++
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/include/wrench/services/compute/cloud/CloudComputeService.h b/include/wrench/services/compute/cloud/CloudComputeService.h
index 6db70b88e7..73f295031c 100644
--- a/include/wrench/services/compute/cloud/CloudComputeService.h
+++ b/include/wrench/services/compute/cloud/CloudComputeService.h
@@ -71,10 +71,10 @@ namespace wrench {
 
     public:
         CloudComputeService(const std::string &hostname,
-                     std::vector<std::string> &execution_hosts,
-                     std::string scratch_space_mount_point,
-                     std::map<std::string, std::string> property_list = {},
-                     std::map<std::string, double> messagepayload_list = {});
+                            std::vector<std::string> &execution_hosts,
+                            std::string scratch_space_mount_point,
+                            std::map<std::string, std::string> property_list = {},
+                            std::map<std::string, double> messagepayload_list = {});
 
         /***********************/
         /** \cond DEVELOPER    */
@@ -95,6 +95,8 @@ namespace wrench {
 
         virtual std::shared_ptr<BareMetalComputeService> startVM(const std::string &vm_name);
 
+        virtual std::shared_ptr<BareMetalComputeService> getVMComputeService(const std::string &vm_name);
+
         virtual void suspendVM(const std::string &vm_name);
 
         virtual void resumeVM(const std::string &vm_name);
diff --git a/src/wrench/services/compute/cloud/CloudComputeService.cpp b/src/wrench/services/compute/cloud/CloudComputeService.cpp
index 149eea1002..e95e86c4fa 100644
--- a/src/wrench/services/compute/cloud/CloudComputeService.cpp
+++ b/src/wrench/services/compute/cloud/CloudComputeService.cpp
@@ -259,6 +259,25 @@ namespace wrench {
         }
     }
 
+    /**
+     * @brief Get the compute service running on a VM, if any
+     *
+     * @param vm_name: the name of the VM
+     *
+     * @return A BareMetalComputeService that runs on the VM, or nullptr if none
+     *
+     * @throw WorkflowExecutionException
+     * @throw std::invalid_argument
+     */
+    std::shared_ptr<BareMetalComputeService> CloudComputeService::getVMComputeService(const std::string &vm_name) {
+
+        if (this->vm_list.find(vm_name) == this->vm_list.end()) {
+            throw std::invalid_argument("CloudComputeService::startVM(): Unknown VM name '" + vm_name + "'");
+        }
+        return this->vm_list.at(vm_name).second;
+    }
+
+
     /**
      * @brief Suspend a running VM
      *
diff --git a/test/compute_services/VirtualizedClusterService/VirtualizedClusterServiceTest.cpp b/test/compute_services/VirtualizedClusterService/VirtualizedClusterServiceTest.cpp
index 8afe0a09df..14e788e143 100644
--- a/test/compute_services/VirtualizedClusterService/VirtualizedClusterServiceTest.cpp
+++ b/test/compute_services/VirtualizedClusterService/VirtualizedClusterServiceTest.cpp
@@ -353,9 +353,20 @@ class CloudStandardJobTestWMS : public wrench::WMS {
 
         }
 
+        // Check that we cannot get the CS back
+        if (cs->getVMComputeService(vm_name) != nullptr) {
+            throw std::runtime_error("A non-started VM should have a nullptr compute service");
+        }
+
         // Start the VM
         auto vm_cs = cs->startVM(vm_name);
 
+        // Check that we can get the CS back
+        auto vm_cs_should_be_same = cs->getVMComputeService(vm_name);
+        if (vm_cs != vm_cs_should_be_same) {
+            throw std::runtime_error("It should be possible to get the computer service of a started VM");
+        }
+
         // Check the state
         if (not cs->isVMRunning(vm_name)) {
             throw std::runtime_error("A just started VM should be running");