Adds get_object test when id_field is uuid.

Author: Florent

elasticutils/contrib/django/tests/__init__.py

@@ -1,6 +1,7 @@
 # We need to put these in a separate module so they're easy to import
 # on a test-by-test basis so that we can skip django-requiring tests
 # if django isn't installed.
+from uuid import UUID
 
 
 from elasticutils.contrib.django import MappingType, Indexable
@@ -22,11 +23,17 @@ class SearchQuerySet(object):
     # Yes. This is kind of crazy, but ... whatever.
     def __init__(self, model):
         self.model = model
+        self.id_field = model.id_field
         self.steps = []
 
-    def get(self, pk):
-        pk = int(pk)
-        return [m for m in _model_cache if m.id == pk][0]
+    def get(self, pk=None, uuid=None):
+        if pk:
+            pk = int(pk)
+            return [m for m in _model_cache if m.id == pk][0]
+        if uuid:
+            uuid = UUID(uuid)
+            return [m for m in _model_cache if m.uuid == uuid][0]
+        return []
 
     def filter(self, id__in=None):
         self.steps.append(('filter', id__in))
@@ -47,7 +54,8 @@ def __iter__(self):
 
         for mem in self.steps:
             if mem[0] == 'filter':
-                objs = [obj for obj in objs if obj.id in mem[1]]
+                objs = [obj for obj in objs
+                        if getattr(obj, self.id_field) in mem[1]]
             elif mem[0] == 'order_by':
                 order_by_field = mem[1][0]
             elif mem[0] == 'values_list':
@@ -58,16 +66,16 @@ def __iter__(self):
 
         if values_list:
             # Note: Hard-coded to just id and flat
-            objs = [obj.id for obj in objs]
+            objs = [getattr(obj, self.id_field) for obj in objs]
         return iter(objs)
 
 
 class Manager(object):
     def get_query_set(self):
         return SearchQuerySet(self)
 
-    def get(self, pk):
-        return self.get_query_set().get(pk)
+    def get(self, pk=None, uuid=None):
+        return self.get_query_set().get(pk=pk, uuid=uuid)
 
     def filter(self, *args, **kwargs):
         return self.get_query_set().filter(*args, **kwargs)
@@ -84,6 +92,7 @@ class FakeModel(object):
     objects = Manager()
 
     def __init__(self, **kw):
+        self.objects.id_field = kw.pop('id_field', 'id')
         self._doc = kw
         for key in kw:
             setattr(self, key, kw[key])
@@ -102,3 +111,6 @@ def extract_document(cls, obj_id, obj=None):
                              'what to do with these args.')
 
         return obj._doc
+
+class FakeDjangoWithUuidMappingType(FakeDjangoMappingType):
+    id_field = 'uuid'

elasticutils/contrib/django/tests/test_models.py

@@ -1,8 +1,11 @@
+import uuid
+
 from nose.tools import eq_
 
 from elasticutils.contrib.django import S, get_es
 from elasticutils.contrib.django.tests import (
-    FakeDjangoMappingType, FakeModel, reset_model_cache)
+    FakeDjangoMappingType, FakeDjangoWithUuidMappingType, FakeModel,
+    reset_model_cache)
 from elasticutils.contrib.django.estestcase import ESTestCase
 
 
@@ -20,12 +23,13 @@ def tearDown(self):
         IndexableTest.cleanup_index(FakeDjangoMappingType.get_index())
         reset_model_cache()
 
-    def persist_data(self, data):
+    def persist_data(self, data, id_field='id'):
         for doc in data:
             FakeModel(**doc)
 
             # Index the document with .index()
-            FakeDjangoMappingType.index(doc, id_=doc['id'])
+            FakeDjangoMappingType.index({k:str(v) for k,v in doc.iteritems()},
+                                        id_=str(doc[id_field]))
 
         self.refresh(FakeDjangoMappingType.get_index())
 
@@ -51,6 +55,17 @@ def test_get_object(self):
         obj = s[0]
         eq_(obj.object.id, 1)
 
+    def test_get_object_with_custom_pk(self):
+        data = [
+            {'uuid': uuid.uuid4(), 'name': 'odin skullcrusher'},
+            {'uuid': uuid.uuid4(), 'name': 'olaf bloodbiter'},
+        ]
+        self.persist_data(data, id_field='uuid')
+
+        s = S(FakeDjangoWithUuidMappingType).query(name__prefix='odin')
+        obj = s[0]
+        eq_(obj.object.uuid, data[0]['uuid'])
+
     def test_get_indexable(self):
         self.persist_data([
                 {'id': 1, 'name': 'odin skullcrusher'},