1 import roslib; roslib.load_manifest('mongo_ros')
   2 import rospy
   3 import mongo_ros as mr
   4 import unittest
   5 import math
   6 import geometry_msgs.msg as gm
   7 
   8 def make_pose(x, y, th):
   9     return gm.Pose(gm.Point(x, y, 0),
  10             gm.Quaternion(0, 0, math.sin(th/2), math.cos(th/2)))
  11     
  12 def make_metadata(p, str):
  13     return {'x': p.position.x, 'y': p.position.y, 'name': str}
  14 
  15 def to_tuple(p):
  16     return (p.position.x, p.position.y, p.position.z, p.orientation.x,
  17             p.orientation.y, p.orientation.z, p.orientation.w)
  18 
  19 def eq_poses(p1, p2):
  20     return to_tuple(p1)==to_tuple(p2)
  21             
  22 
  23 class TestMongoRospy(unittest.TestCase):
  24 
  25     def test_basic(self):
  26 
  27         # Set up collection
  28         coll = mr.MessageCollection("my_db", "poses", gm.Pose)
  29         p1 = make_pose(24, 42, 0)
  30         p2 = make_pose(10, 532, 3)
  31         p3 = make_pose(53, 22, 5)
  32         p4 = make_pose(22, -5, 33)
  33 
  34         # Insert pose objects with accompanying string metadata
  35         coll.insert(p1, make_metadata(p1, "bar"))
  36         coll.insert(p2, make_metadata(p2, "baz"))
  37         coll.insert(p3, make_metadata(p3, "qux"))
  38         coll.insert(p1, make_metadata(p1, "oof"))
  39         coll.insert(p4, make_metadata(p4, "ooof"))
  40 
  41         # Query poses s.t x < 40, y > ), in descending order of name
  42         results = coll.query({'x': {'$lt': 40}, 'y': {'$gt': 0}},\
  43                 sort_by='name', ascending=False)
  44 
  45         # Turn list of pairs into pair of lists
  46         poses, metadata = zip(*list(results))
  47         
  48         self.assertEqual(len(poses), 3)
  49         self.assertTrue(eq_poses(p1, poses[0]))
  50         self.assertTrue(eq_poses(p2, poses[1]))
  51         self.assertTrue(eq_poses(p1, poses[2]))
  52 
  53         self.assertEqual(metadata[0]['name'], 'oof')
  54         self.assertEqual(metadata[1]['name'], 'baz')
  55         self.assertEqual(metadata[2]['name'], 'bar')
  56 
  57         # Update some messages
  58         coll.update(metadata[0], metadata={'name': 'bat'})
  59         coll.update(metadata[2], msg=p2)
  60         res2 = coll.query({'y': {'$gt': 40}}, sort_by='name')
  61         poses, metadata = zip(*list(res2))
  62 
  63         self.assertEqual(metadata[0]['name'], 'bar')
  64         self.assertEqual(metadata[1]['name'], 'bat')
  65         self.assertEqual(metadata[2]['name'], 'baz')
  66         self.assertTrue(eq_poses(p2, poses[0]))
  67         self.assertTrue(eq_poses(p1, poses[1]))
  68         self.assertTrue(eq_poses(p2, poses[2]))
  69         
  70 
  71         # Remove entries s.t. y<30
  72         self.assertEqual(5, coll.count())
  73         self.assertEqual(2, coll.remove({'y': {'$lt': 30}}))
  74         self.assertEqual(3, coll.count())
  75 
  76         # Test find_one
  77         self.assertTrue(coll.find_one({'y': {'$gt': 30}}))
  78         self.assertFalse(coll.find_one({'y': {'$lt': 30}}))
  79         
  80 
  81 if __name__ == "__main__":
  82     rospy.init_node('test_mongo_rospy')
  83     import rostest
  84     rostest.rosrun('mongo_ros', 'test_mongo_rospy', TestMongoRospy)