From 55a14ca8b6d298b9ddc0e2f5fa4c9cca5871f30c Mon Sep 17 00:00:00 2001
From: Anting Shen <antingshen@berkeley.edu>
Date: Tue, 18 Nov 2014 20:56:25 -0800
Subject: [PATCH] fix update/step assumption that a robot has been added & add
 test to prevent these assumptions

---
 lfd/environment/simulation.py |  6 ++++--
 test/test_simulation.py       | 17 +++++++++--------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/lfd/environment/simulation.py b/lfd/environment/simulation.py
index 2fa4fc1..239a073 100644
--- a/lfd/environment/simulation.py
+++ b/lfd/environment/simulation.py
@@ -195,11 +195,13 @@ def set_state(self, sim_state):
         self.update()
 
     def update(self):
-        self.bt_robot.UpdateBullet()
+        if self.robot:
+            self.bt_robot.UpdateBullet()
         self._update_rave()
     
     def step(self):
-        self.bt_robot.UpdateBullet()
+        if self.robot:
+            self.bt_robot.UpdateBullet()
         self.bt_env.Step(.01, 200, .005)
         self._update_rave()
 
diff --git a/test/test_simulation.py b/test/test_simulation.py
index 1141cfe..2c6527c 100644
--- a/test/test_simulation.py
+++ b/test/test_simulation.py
@@ -35,7 +35,6 @@ def setUp(self):
         cyl_pos1 = np.r_[.6, -helix_radius, table_height + .35]
         
         sim_objs = []
-        sim_objs.append(XmlSimulationObject("robots/pr2-beta-static.zae", dynamic=False))
         sim_objs.append(BoxSimulationObject("table", [1, 0, table_height-.1], [.85, .85, .1], dynamic=False))
         sim_objs.append(RopeSimulationObject("rope", init_rope_nodes, rope_params))
         sim_objs.append(CylinderSimulationObject("cyl0", cyl_pos0, cyl_radius, cyl_height, dynamic=True))
@@ -43,16 +42,18 @@ def setUp(self):
         
         self.sim = DynamicSimulation()
         self.sim.add_objects(sim_objs)
-        self.sim.robot.SetDOFValues([0.25], [self.sim.robot.GetJoint('torso_lift_joint').GetJointIndex()])
-        sim_util.reset_arms_to_side(self.sim)
-        
+
         # rotate cylinders by 90 deg
         for i in range(2):
-            bt_cyl = self.sim.bt_env.GetObjectByName('cyl%d'%i)
-            T = openravepy.matrixFromAxisAngle(np.array([np.pi/2,0,0]))
-            T[:3,3] = bt_cyl.GetTransform()[:3,3]
-            bt_cyl.SetTransform(T) # SetTransform needs to be used in the Bullet object, not the openrave body
+            bt_cyl = self.sim.bt_env.GetObjectByName('cyl%d' % i)
+            T = openravepy.matrixFromAxisAngle(np.array([np.pi/2, 0, 0]))
+            T[:3,3] = bt_cyl.GetTransform()[:3, 3]
+            bt_cyl.SetTransform(T)  # SetTransform needs to be used in the Bullet object, not the openrave body
         self.sim.update()
+
+        self.sim.add_objects([XmlSimulationObject("robots/pr2-beta-static.zae", dynamic=False)])
+        self.sim.robot.SetDOFValues([0.25], [self.sim.robot.GetJoint('torso_lift_joint').GetJointIndex()])
+        sim_util.reset_arms_to_side(self.sim)
     
     def test_reproducibility(self):
         sim_state0 = self.sim.get_state()