diff --git a/tests/queryset/geo.py b/tests/queryset/geo.py index 0edde4db..b6e602fc 100644 --- a/tests/queryset/geo.py +++ b/tests/queryset/geo.py @@ -1,13 +1,9 @@ import datetime import unittest -from nose.plugins.skip import SkipTest -from pymongo.errors import OperationFailure - from mongoengine import * -from mongoengine.connection import get_connection -from tests.utils import MongoDBTestCase, get_mongodb_version +from tests.utils import MongoDBTestCase, skip_older_mongodb __all__ = ("GeoQueriesTest",) @@ -15,94 +11,128 @@ __all__ = ("GeoQueriesTest",) class GeoQueriesTest(MongoDBTestCase): - def setUp(self): - connect(db='mongoenginetest') - - def test_geospatial_operators(self): - """Ensure that geospatial queries are working. - """ + def _create_event_data(self, point_field_class=GeoPointField): + """Create some sample data re-used in many of the tests below.""" class Event(Document): title = StringField() date = DateTimeField() - location = GeoPointField() + location = point_field_class() def __unicode__(self): return self.title + self.Event = Event + Event.drop_collection() - event1 = Event(title="Coltrane Motion @ Double Door", - date=datetime.datetime.now() - datetime.timedelta(days=1), - location=[-87.677137, 41.909889]).save() - event2 = Event(title="Coltrane Motion @ Bottom of the Hill", - date=datetime.datetime.now() - datetime.timedelta(days=10), - location=[-122.4194155, 37.7749295]).save() - event3 = Event(title="Coltrane Motion @ Empty Bottle", - date=datetime.datetime.now(), - location=[-87.686638, 41.900474]).save() + event1 = Event.objects.create( + title="Coltrane Motion @ Double Door", + date=datetime.datetime.now() - datetime.timedelta(days=1), + location=[-87.677137, 41.909889]) + event2 = Event.objects.create( + title="Coltrane Motion @ Bottom of the Hill", + date=datetime.datetime.now() - datetime.timedelta(days=10), + location=[-122.4194155, 37.7749295]) + event3 = Event.objects.create( + title="Coltrane Motion @ Empty Bottle", + date=datetime.datetime.now(), + location=[-87.686638, 41.900474]) + + return event1, event2, event3 + + def test_near(self): + """Make sure the "near" operator works.""" + event1, event2, event3 = self._create_event_data() # find all events "near" pitchfork office, chicago. # note that "near" will show the san francisco event, too, # although it sorts to last. - events = Event.objects(location__near=[-87.67892, 41.9120459]) + events = self.Event.objects(location__near=[-87.67892, 41.9120459]) self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event1, event3, event2]) + # ensure ordering is respected by "near" + events = self.Event.objects(location__near=[-87.67892, 41.9120459]) + events = events.order_by("-date") + self.assertEqual(events.count(), 3) + self.assertEqual(list(events), [event3, event1, event2]) + + def test_near_and_max_distance(self): + """Ensure the "max_distance" operator works alongside the "near" + operator. + """ + event1, event2, event3 = self._create_event_data() + + # find events within 10 degrees of san francisco + point = [-122.415579, 37.7566023] + events = self.Event.objects(location__near=point, + location__max_distance=10) + self.assertEqual(events.count(), 1) + self.assertEqual(events[0], event2) + + # $minDistance was only added in MongoDB v2.6, skip for older versions + @skip_older_mongodb + def test_near_and_min_distance(self): + """Ensure the "min_distance" operator works alongside the "near" + operator. + """ + event1, event2, event3 = self._create_event_data() + + # find events at least 10 degrees away of san francisco + point = [-122.415579, 37.7566023] + events = self.Event.objects(location__near=point, + location__min_distance=10) + self.assertEqual(events.count(), 2) + + def test_within_distance(self): + """Make sure the "within_distance" operator works.""" + event1, event2, event3 = self._create_event_data() + # find events within 5 degrees of pitchfork office, chicago point_and_distance = [[-87.67892, 41.9120459], 5] - events = Event.objects(location__within_distance=point_and_distance) + events = self.Event.objects( + location__within_distance=point_and_distance) self.assertEqual(events.count(), 2) events = list(events) self.assertTrue(event2 not in events) self.assertTrue(event1 in events) self.assertTrue(event3 in events) - # ensure ordering is respected by "near" - events = Event.objects(location__near=[-87.67892, 41.9120459]) - events = events.order_by("-date") - self.assertEqual(events.count(), 3) - self.assertEqual(list(events), [event3, event1, event2]) - - # find events within 10 degrees of san francisco - point = [-122.415579, 37.7566023] - events = Event.objects(location__near=point, location__max_distance=10) - self.assertEqual(events.count(), 1) - self.assertEqual(events[0], event2) - - # find events at least 10 degrees away of san francisco - point = [-122.415579, 37.7566023] - events = Event.objects(location__near=point, location__min_distance=10) - # The following real test passes on MongoDB 3 but minDistance seems - # buggy on older MongoDB versions - if get_connection().server_info()['versionArray'][0] > 2: - self.assertEqual(events.count(), 2) - else: - self.assertTrue(events.count() >= 2) - # find events within 10 degrees of san francisco point_and_distance = [[-122.415579, 37.7566023], 10] - events = Event.objects(location__within_distance=point_and_distance) + events = self.Event.objects( + location__within_distance=point_and_distance) self.assertEqual(events.count(), 1) self.assertEqual(events[0], event2) # find events within 1 degree of greenpoint, broolyn, nyc, ny point_and_distance = [[-73.9509714, 40.7237134], 1] - events = Event.objects(location__within_distance=point_and_distance) + events = self.Event.objects( + location__within_distance=point_and_distance) self.assertEqual(events.count(), 0) # ensure ordering is respected by "within_distance" point_and_distance = [[-87.67892, 41.9120459], 10] - events = Event.objects(location__within_distance=point_and_distance) + events = self.Event.objects( + location__within_distance=point_and_distance) events = events.order_by("-date") self.assertEqual(events.count(), 2) self.assertEqual(events[0], event3) + def test_within_box(self): + """Ensure the "within_box" operator works.""" + event1, event2, event3 = self._create_event_data() + # check that within_box works box = [(-125.0, 35.0), (-100.0, 40.0)] - events = Event.objects(location__within_box=box) + events = self.Event.objects(location__within_box=box) self.assertEqual(events.count(), 1) self.assertEqual(events[0].id, event2.id) + def test_within_polygon(self): + """Ensure the "within_polygon" operator works.""" + event1, event2, event3 = self._create_event_data() + polygon = [ (-87.694445, 41.912114), (-87.69084, 41.919395), @@ -110,7 +140,7 @@ class GeoQueriesTest(MongoDBTestCase): (-87.654276, 41.911731), (-87.656164, 41.898061), ] - events = Event.objects(location__within_polygon=polygon) + events = self.Event.objects(location__within_polygon=polygon) self.assertEqual(events.count(), 1) self.assertEqual(events[0].id, event1.id) @@ -119,13 +149,150 @@ class GeoQueriesTest(MongoDBTestCase): (-1.225891, 52.792797), (-4.40094, 53.389881) ] - events = Event.objects(location__within_polygon=polygon2) + events = self.Event.objects(location__within_polygon=polygon2) self.assertEqual(events.count(), 0) - def test_geo_spatial_embedded(self): + def test_2dsphere_near(self): + """Make sure the "near" operator works with a PointField, which + corresponds to a 2dsphere index. + """ + event1, event2, event3 = self._create_event_data( + point_field_class=PointField + ) + # find all events "near" pitchfork office, chicago. + # note that "near" will show the san francisco event, too, + # although it sorts to last. + events = self.Event.objects(location__near=[-87.67892, 41.9120459]) + self.assertEqual(events.count(), 3) + self.assertEqual(list(events), [event1, event3, event2]) + + # ensure ordering is respected by "near" + events = self.Event.objects(location__near=[-87.67892, 41.9120459]) + events = events.order_by("-date") + self.assertEqual(events.count(), 3) + self.assertEqual(list(events), [event3, event1, event2]) + + def test_2dsphere_near_and_max_distance(self): + """Ensure the "max_distance" operator works alongside the "near" + operator with a 2dsphere index. + """ + event1, event2, event3 = self._create_event_data( + point_field_class=PointField + ) + + # find events within 10km of san francisco + point = [-122.415579, 37.7566023] + events = self.Event.objects(location__near=point, + location__max_distance=10000) + self.assertEqual(events.count(), 1) + self.assertEqual(events[0], event2) + + # find events within 1km of greenpoint, broolyn, nyc, ny + events = self.Event.objects(location__near=[-73.9509714, 40.7237134], + location__max_distance=1000) + self.assertEqual(events.count(), 0) + + # ensure ordering is respected by "near" + events = self.Event.objects( + location__near=[-87.67892, 41.9120459], + location__max_distance=10000 + ).order_by("-date") + self.assertEqual(events.count(), 2) + self.assertEqual(events[0], event3) + + def test_2dsphere_geo_within_box(self): + """Ensure the "geo_within_box" operator works with a 2dsphere + index. + """ + event1, event2, event3 = self._create_event_data( + point_field_class=PointField + ) + + # check that within_box works + box = [(-125.0, 35.0), (-100.0, 40.0)] + events = self.Event.objects(location__geo_within_box=box) + self.assertEqual(events.count(), 1) + self.assertEqual(events[0].id, event2.id) + + def test_2dsphere_geo_within_polygon(self): + """Ensure the "geo_within_polygon" operator works with a + 2dsphere index. + """ + event1, event2, event3 = self._create_event_data( + point_field_class=PointField + ) + + polygon = [ + (-87.694445, 41.912114), + (-87.69084, 41.919395), + (-87.681742, 41.927186), + (-87.654276, 41.911731), + (-87.656164, 41.898061), + ] + events = self.Event.objects(location__geo_within_polygon=polygon) + self.assertEqual(events.count(), 1) + self.assertEqual(events[0].id, event1.id) + + polygon2 = [ + (-1.742249, 54.033586), + (-1.225891, 52.792797), + (-4.40094, 53.389881) + ] + events = self.Event.objects(location__geo_within_polygon=polygon2) + self.assertEqual(events.count(), 0) + + # $minDistance was only added in MongoDB v2.6, skip for older versions + @skip_older_mongodb + def test_2dsphere_near_and_min_max_distance(self): + """Ensure "min_distace" and "max_distance" operators work well + together with the "near" operator in a 2dsphere index. + """ + event1, event2, event3 = self._create_event_data( + point_field_class=PointField + ) + + # ensure min_distance and max_distance combine well + events = self.Event.objects( + location__near=[-87.67892, 41.9120459], + location__min_distance=1000, + location__max_distance=10000 + ).order_by("-date") + self.assertEqual(events.count(), 1) + self.assertEqual(events[0], event3) + + # ensure ordering is respected by "near" with "min_distance" + events = self.Event.objects( + location__near=[-87.67892, 41.9120459], + location__min_distance=10000 + ).order_by("-date") + self.assertEqual(events.count(), 1) + self.assertEqual(events[0], event2) + + def test_2dsphere_geo_within_center(self): + """Make sure the "geo_within_center" operator works with a + 2dsphere index. + """ + event1, event2, event3 = self._create_event_data( + point_field_class=PointField + ) + + # find events within 5 degrees of pitchfork office, chicago + point_and_distance = [[-87.67892, 41.9120459], 2] + events = self.Event.objects( + location__geo_within_center=point_and_distance) + self.assertEqual(events.count(), 2) + events = list(events) + self.assertTrue(event2 not in events) + self.assertTrue(event1 in events) + self.assertTrue(event3 in events) + + def _test_embedded(self, point_field_class): + """Helper test method ensuring given point field class works + well in an embedded document. + """ class Venue(EmbeddedDocument): - location = GeoPointField() + location = point_field_class() name = StringField() class Event(Document): @@ -151,14 +318,18 @@ class GeoQueriesTest(MongoDBTestCase): self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event1, event3, event2]) - def test_spherical_geospatial_operators(self): - """Ensure that spherical geospatial queries are working - """ - # Needs MongoDB > 2.6.4 https://jira.mongodb.org/browse/SERVER-14039 - mongodb_version = get_mongodb_version() - if mongodb_version[0] <= 2 and mongodb_version[1] <= 6 and mongodb_version[2] <= 4: - raise SkipTest("Need MongoDB version 2.6.4+") + def test_geo_spatial_embedded(self): + """Make sure GeoPointField works properly in an embedded document.""" + self._test_embedded(point_field_class=GeoPointField) + def test_2dsphere_point_embedded(self): + """Make sure PointField works properly in an embedded document.""" + self._test_embedded(point_field_class=PointField) + + # Needs MongoDB > 2.6.4 https://jira.mongodb.org/browse/SERVER-14039 + @skip_older_mongodb + def test_spherical_geospatial_operators(self): + """Ensure that spherical geospatial queries are working.""" class Point(Document): location = GeoPointField() @@ -178,7 +349,10 @@ class GeoQueriesTest(MongoDBTestCase): # Same behavior for _within_spherical_distance points = Point.objects( - location__within_spherical_distance=[[-122, 37.5], 60 / earth_radius] + location__within_spherical_distance=[ + [-122, 37.5], + 60 / earth_radius + ] ) self.assertEqual(points.count(), 2) @@ -195,14 +369,9 @@ class GeoQueriesTest(MongoDBTestCase): # Test query works with min_distance, being farer from one point points = Point.objects(location__near_sphere=[-122, 37.8], location__min_distance=60 / earth_radius) - # The following real test passes on MongoDB 3 but minDistance seems - # buggy on older MongoDB versions - if get_connection().server_info()['versionArray'][0] > 2: - self.assertEqual(points.count(), 1) - far_point = points.first() - self.assertNotEqual(close_point, far_point) - else: - self.assertTrue(points.count() >= 1) + self.assertEqual(points.count(), 1) + far_point = points.first() + self.assertNotEqual(close_point, far_point) # Finds both points, but orders the north point first because it's # closer to the reference point to the north. @@ -221,146 +390,15 @@ class GeoQueriesTest(MongoDBTestCase): # Finds only one point because only the first point is within 60km of # the reference point to the south. points = Point.objects( - location__within_spherical_distance=[[-122, 36.5], 60/earth_radius]) + location__within_spherical_distance=[ + [-122, 36.5], + 60 / earth_radius + ] + ) self.assertEqual(points.count(), 1) self.assertEqual(points[0].id, south_point.id) - def test_2dsphere_point(self): - class Event(Document): - title = StringField() - date = DateTimeField() - location = PointField() - - def __unicode__(self): - return self.title - - Event.drop_collection() - - event1 = Event(title="Coltrane Motion @ Double Door", - date=datetime.datetime.now() - datetime.timedelta(days=1), - location=[-87.677137, 41.909889]).save() - event2 = Event(title="Coltrane Motion @ Bottom of the Hill", - date=datetime.datetime.now() - datetime.timedelta(days=10), - location=[-122.4194155, 37.7749295]).save() - event3 = Event(title="Coltrane Motion @ Empty Bottle", - date=datetime.datetime.now(), - location=[-87.686638, 41.900474]).save() - - # find all events "near" pitchfork office, chicago. - # note that "near" will show the san francisco event, too, - # although it sorts to last. - events = Event.objects(location__near=[-87.67892, 41.9120459]) - self.assertEqual(events.count(), 3) - self.assertEqual(list(events), [event1, event3, event2]) - - # find events within 5 degrees of pitchfork office, chicago - point_and_distance = [[-87.67892, 41.9120459], 2] - events = Event.objects(location__geo_within_center=point_and_distance) - self.assertEqual(events.count(), 2) - events = list(events) - self.assertTrue(event2 not in events) - self.assertTrue(event1 in events) - self.assertTrue(event3 in events) - - # ensure ordering is respected by "near" - events = Event.objects(location__near=[-87.67892, 41.9120459]) - events = events.order_by("-date") - self.assertEqual(events.count(), 3) - self.assertEqual(list(events), [event3, event1, event2]) - - # find events within 10km of san francisco - point = [-122.415579, 37.7566023] - events = Event.objects(location__near=point, - location__max_distance=10000) - self.assertEqual(events.count(), 1) - self.assertEqual(events[0], event2) - - # find events within 1km of greenpoint, broolyn, nyc, ny - events = Event.objects(location__near=[-73.9509714, 40.7237134], - location__max_distance=1000) - self.assertEqual(events.count(), 0) - - # ensure ordering is respected by "near" - events = Event.objects(location__near=[-87.67892, 41.9120459], - location__max_distance=10000).order_by("-date") - self.assertEqual(events.count(), 2) - self.assertEqual(events[0], event3) - - # check that within_box works - box = [(-125.0, 35.0), (-100.0, 40.0)] - events = Event.objects(location__geo_within_box=box) - self.assertEqual(events.count(), 1) - self.assertEqual(events[0].id, event2.id) - - polygon = [ - (-87.694445, 41.912114), - (-87.69084, 41.919395), - (-87.681742, 41.927186), - (-87.654276, 41.911731), - (-87.656164, 41.898061), - ] - events = Event.objects(location__geo_within_polygon=polygon) - self.assertEqual(events.count(), 1) - self.assertEqual(events[0].id, event1.id) - - polygon2 = [ - (-1.742249, 54.033586), - (-1.225891, 52.792797), - (-4.40094, 53.389881) - ] - events = Event.objects(location__geo_within_polygon=polygon2) - self.assertEqual(events.count(), 0) - - # $minDistance was only added in MongoDB v2.6, skip for older versions - mongodb_ver = get_mongodb_version() - if mongodb_ver[0] == 2 and mongodb_ver[1] < 6: - raise SkipTest('Need MongoDB v2.6+') - - # ensure min_distance and max_distance combine well - events = Event.objects(location__near=[-87.67892, 41.9120459], - location__min_distance=1000, - location__max_distance=10000).order_by("-date") - self.assertEqual(events.count(), 1) - self.assertEqual(events[0], event3) - - # ensure ordering is respected by "near" - events = Event.objects(location__near=[-87.67892, 41.9120459], - # location__min_distance=10000 - location__min_distance=10000).order_by("-date") - self.assertEqual(events.count(), 1) - self.assertEqual(events[0], event2) - - def test_2dsphere_point_embedded(self): - - class Venue(EmbeddedDocument): - location = GeoPointField() - name = StringField() - - class Event(Document): - title = StringField() - venue = EmbeddedDocumentField(Venue) - - Event.drop_collection() - - venue1 = Venue(name="The Rock", location=[-87.677137, 41.909889]) - venue2 = Venue(name="The Bridge", location=[-122.4194155, 37.7749295]) - - event1 = Event(title="Coltrane Motion @ Double Door", - venue=venue1).save() - event2 = Event(title="Coltrane Motion @ Bottom of the Hill", - venue=venue2).save() - event3 = Event(title="Coltrane Motion @ Empty Bottle", - venue=venue1).save() - - # find all events "near" pitchfork office, chicago. - # note that "near" will show the san francisco event, too, - # although it sorts to last. - events = Event.objects(venue__location__near=[-87.67892, 41.9120459]) - self.assertEqual(events.count(), 3) - self.assertEqual(list(events), [event1, event3, event2]) - def test_linestring(self): - class Road(Document): name = StringField() line = LineStringField() @@ -416,7 +454,6 @@ class GeoQueriesTest(MongoDBTestCase): self.assertEqual(1, roads) def test_polygon(self): - class Road(Document): name = StringField() poly = PolygonField() @@ -513,5 +550,6 @@ class GeoQueriesTest(MongoDBTestCase): loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["poly"], {"type": "Polygon", "coordinates": [[[40, 4], [40, 6], [41, 6], [40, 4]]]}) + if __name__ == '__main__': unittest.main()