From a9563f5b0e16d1495a2b3b3b3937db12c5e43008 Mon Sep 17 00:00:00 2001
From: Ian Chen <ichen@osrfoundation.org>
Date: Tue, 15 Mar 2022 08:57:51 -0700
Subject: [PATCH] Add gazebo Entity id to rendering sensor's user data (#1381)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds the gazebo::Entity id to a rendering::Sensor's UserData object. The main use case is so that we can get back the corresponding gazebo Entity in the rendering thread when processing sensors.

Signed-off-by: Ian Chen <ichen@osrfoundation.org>

Co-authored-by: Nate Koenig <nkoenig@users.noreply.github.com>
Co-authored-by: Alejandro Hernández Cordero <ahcorde@gmail.com>
---
 src/rendering/SceneManager.cc      |  3 ++
 test/integration/CMakeLists.txt    |  6 +++
 test/integration/sensors_system.cc | 63 ++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/src/rendering/SceneManager.cc b/src/rendering/SceneManager.cc
index 613210ce62..ac6720231d 100644
--- a/src/rendering/SceneManager.cc
+++ b/src/rendering/SceneManager.cc
@@ -1776,6 +1776,9 @@ bool SceneManager::AddSensor(Entity _gazeboId, const std::string &_sensorName,
     return false;
   }
 
+  // \todo(anyone) change to uint64_t once UserData supports this type
+  sensor->SetUserData("gazebo-entity", static_cast<int>(_gazeboId));
+
   if (parent)
   {
     sensor->RemoveParent();
diff --git a/test/integration/CMakeLists.txt b/test/integration/CMakeLists.txt
index 4a8bee7003..ab9f4cec3a 100644
--- a/test/integration/CMakeLists.txt
+++ b/test/integration/CMakeLists.txt
@@ -139,3 +139,9 @@ set_tests_properties(INTEGRATION_tracked_vehicle_system PROPERTIES TIMEOUT 300)
 if(TARGET INTEGRATION_examples_build)
   set_tests_properties(INTEGRATION_examples_build PROPERTIES TIMEOUT 320)
 endif()
+
+if(VALID_DISPLAY AND VALID_DRI_DISPLAY)
+  target_link_libraries(INTEGRATION_sensors_system
+    ignition-rendering${IGN_RENDERING_VER}::ignition-rendering${IGN_RENDERING_VER}
+  )
+endif()
diff --git a/test/integration/sensors_system.cc b/test/integration/sensors_system.cc
index 6165ee9620..f640add611 100644
--- a/test/integration/sensors_system.cc
+++ b/test/integration/sensors_system.cc
@@ -18,6 +18,7 @@
 #include <gtest/gtest.h>
 
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #include <sdf/Model.hh>
@@ -27,6 +28,11 @@
 #include <ignition/transport/Node.hh>
 #include <ignition/utilities/ExtraTestMacros.hh>
 
+#include <ignition/rendering/Camera.hh>
+#include <ignition/rendering/RenderEngine.hh>
+#include <ignition/rendering/RenderingIface.hh>
+#include <ignition/rendering/Scene.hh>
+
 #include "ignition/gazebo/EntityComponentManager.hh"
 #include "ignition/gazebo/EventManager.hh"
 #include "ignition/gazebo/SdfEntityCreator.hh"
@@ -37,8 +43,11 @@
 
 #include "ignition/gazebo/components/Model.hh"
 #include "ignition/gazebo/components/Name.hh"
+#include "ignition/gazebo/components/Sensor.hh"
 #include "ignition/gazebo/components/World.hh"
 
+#include "ignition/gazebo/rendering/Events.hh"
+
 #include "plugins/MockSystem.hh"
 #include "../helpers/EnvTestFixture.hh"
 
@@ -46,6 +55,45 @@ using namespace ignition;
 using namespace std::chrono_literals;
 namespace components = ignition::gazebo::components;
 
+std::unordered_set<gazebo::Entity> g_sensorEntityIds;
+rendering::ScenePtr g_scene;
+
+/////////////////////////////////////////////////
+void OnPostRender()
+{
+  if (!g_scene)
+  {
+    g_scene = rendering::sceneFromFirstRenderEngine();
+  }
+  ASSERT_TRUE(g_scene);
+
+  EXPECT_LT(0u, g_scene->SensorCount());
+  for (unsigned int i = 0; i < g_scene->SensorCount(); ++i)
+  {
+    auto sensor = g_scene->SensorByIndex(i);
+    ASSERT_TRUE(sensor);
+    EXPECT_TRUE(sensor->HasUserData("gazebo-entity"));
+    auto variant = sensor->UserData("gazebo-entity");
+
+    // todo(anyone) change int to uint64_t once user data supports it
+    const int *value = std::get_if<int>(&variant);
+    ASSERT_TRUE(value);
+    g_sensorEntityIds.insert(*value);
+  }
+}
+
+/////////////////////////////////////////////////
+void testSensorEntityIds(const gazebo::EntityComponentManager &_ecm,
+    const std::unordered_set<gazebo::Entity> &_ids)
+{
+  EXPECT_FALSE(_ids.empty());
+  for (const auto & id : _ids)
+  {
+    auto sensorComp = _ecm.Component<gazebo::components::Sensor>(id);
+    EXPECT_TRUE(sensorComp);
+  }
+}
+
 //////////////////////////////////////////////////
 class SensorsFixture : public InternalFixture<InternalFixture<::testing::Test>>
 {
@@ -127,6 +175,8 @@ TEST_F(SensorsFixture, IGN_UTILS_TEST_DISABLED_ON_MAC(HandleRemovedEntities))
 
   gazebo::Server server(serverConfig);
 
+  common::ConnectionPtr postRenderConn;
+
   // A pointer to the ecm. This will be valid once we run the mock system
   gazebo::EntityComponentManager *ecm = nullptr;
   this->mockSystem->preUpdateCallback =
@@ -134,6 +184,15 @@ TEST_F(SensorsFixture, IGN_UTILS_TEST_DISABLED_ON_MAC(HandleRemovedEntities))
     {
       ecm = &_ecm;
     };
+  this->mockSystem->configureCallback =
+    [&](const gazebo::Entity &,
+           const std::shared_ptr<const sdf::Element> &,
+           gazebo::EntityComponentManager &,
+           gazebo::EventManager &_eventMgr)
+    {
+      postRenderConn = _eventMgr.Connect<gazebo::events::PostRender>(
+          std::bind(&::OnPostRender));
+    };
 
   server.AddSystem(this->systemPtr);
   server.Run(true, 50, false);
@@ -141,6 +200,10 @@ TEST_F(SensorsFixture, IGN_UTILS_TEST_DISABLED_ON_MAC(HandleRemovedEntities))
 
   testDefaultTopics();
 
+  testSensorEntityIds(*ecm, g_sensorEntityIds);
+  g_sensorEntityIds.clear();
+  g_scene.reset();
+
   // We won't use the event manager but it's required to create an
   // SdfEntityCreator
   gazebo::EventManager dummyEventMgr;