From 315424c4f0b78af51e8485b270e5890af30ca0e8 Mon Sep 17 00:00:00 2001
From: Steven Peters <scpeters@openrobotics.org>
Date: Tue, 13 Jul 2021 22:12:52 -0700
Subject: [PATCH] Test Gravity APIs in WorldFeatures_TEST

Signed-off-by: Steven Peters <scpeters@openrobotics.org>
---
 dartsim/src/WorldFeatures_TEST.cc | 95 +++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)

diff --git a/dartsim/src/WorldFeatures_TEST.cc b/dartsim/src/WorldFeatures_TEST.cc
index 0de31feaa..b6b510473 100644
--- a/dartsim/src/WorldFeatures_TEST.cc
+++ b/dartsim/src/WorldFeatures_TEST.cc
@@ -35,6 +35,8 @@
 
 struct TestFeatureList : ignition::physics::FeatureList<
     ignition::physics::CollisionDetector,
+    ignition::physics::Gravity,
+    ignition::physics::LinkFrameSemantics,
     ignition::physics::Solver,
     ignition::physics::ForwardStep,
     ignition::physics::sdf::ConstructSdfWorld,
@@ -46,6 +48,29 @@ using namespace ignition;
 using TestEnginePtr = physics::Engine3dPtr<TestFeatureList>;
 using TestWorldPtr = physics::World3dPtr<TestFeatureList>;
 
+// A predicate-formatter for asserting that two vectors are approximately equal.
+class AssertVectorApprox
+{
+  public: explicit AssertVectorApprox(double _tol = 1e-6) : tol(_tol)
+  {
+  }
+
+  public: ::testing::AssertionResult operator()(
+              const char *_mExpr, const char *_nExpr, Eigen::Vector3d _m,
+              Eigen::Vector3d _n)
+  {
+    if (ignition::physics::test::Equal(_m, _n, this->tol))
+      return ::testing::AssertionSuccess();
+
+    return ::testing::AssertionFailure()
+           << _mExpr << " and " << _nExpr << " ([" << _m.transpose()
+           << "] and [" << _n.transpose() << "]"
+           << ") are not equal";
+  }
+
+  private: double tol;
+};
+
 //////////////////////////////////////////////////
 class WorldFeaturesFixture : public ::testing::Test
 {
@@ -98,6 +123,76 @@ TEST_F(WorldFeaturesFixture, CollisionDetector)
   EXPECT_EQ("dart", world->GetCollisionDetector());
 }
 
+//////////////////////////////////////////////////
+TEST_F(WorldFeaturesFixture, Gravity)
+{
+  auto world = LoadWorld(this->engine, TEST_WORLD_DIR "/falling.world");
+  ASSERT_NE(nullptr, world);
+
+  // Check default gravity value
+  AssertVectorApprox vectorPredicate10(1e-10);
+  EXPECT_PRED_FORMAT2(vectorPredicate10, Eigen::Vector3d(0, 0, -9.8),
+                      world->GetGravity());
+
+  auto model = world->GetModel("sphere");
+  ASSERT_NE(nullptr, model);
+
+  auto link = model->GetLink(0);
+  ASSERT_NE(nullptr, link);
+
+  // initial link pose
+  const Eigen::Vector3d initialLinkPosition(0, 0, 2);
+  {
+    auto pos = link->FrameDataRelativeToWorld().pose.translation();
+    EXPECT_PRED_FORMAT2(vectorPredicate10,
+                        initialLinkPosition,
+                        pos);
+  }
+
+  auto linkFrameID = link->GetFrameID();
+
+  // Get default gravity in link frame, which is pitched by pi/4
+  EXPECT_PRED_FORMAT2(vectorPredicate10,
+                      Eigen::Vector3d(6.92964645563, 0, -6.92964645563),
+                      world->GetGravity(linkFrameID));
+
+  // set gravity along X axis of linked frame, which is pitched by pi/4
+  world->SetGravity(Eigen::Vector3d(1.4142135624, 0, 0), linkFrameID);
+
+  EXPECT_PRED_FORMAT2(vectorPredicate10,
+                      Eigen::Vector3d(1, 0, -1),
+                      world->GetGravity());
+
+  // test other SetGravity API
+  // set gravity along Z axis of linked frame, which is pitched by pi/4
+  physics::RelativeForce3d relativeGravity(
+      linkFrameID, Eigen::Vector3d(0, 0, 1.4142135624));
+  world->SetGravity(relativeGravity);
+
+  EXPECT_PRED_FORMAT2(vectorPredicate10,
+                      Eigen::Vector3d(1, 0, 1),
+                      world->GetGravity());
+
+  // Confirm that changed gravity direction affects pose of link
+  ignition::physics::ForwardStep::Input input;
+  ignition::physics::ForwardStep::State state;
+  ignition::physics::ForwardStep::Output output;
+
+  const size_t numSteps = 1000;
+  for (size_t i = 0; i < numSteps; ++i)
+  {
+    world->Step(output, state, input);
+  }
+
+  AssertVectorApprox vectorPredicate3(1e-3);
+  {
+    auto pos = link->FrameDataRelativeToWorld().pose.translation();
+    EXPECT_PRED_FORMAT2(vectorPredicate3,
+                        Eigen::Vector3d(0.5, 0, 2.5),
+                        pos);
+  }
+}
+
 //////////////////////////////////////////////////
 TEST_F(WorldFeaturesFixture, Solver)
 {