From 0d92340f2b1cab932f14e376dc776a7c11e3e42f Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Wed, 4 Dec 2013 15:54:30 +0200
Subject: [PATCH] Add tests for generic relationship batch fetching

---
 .../batch_fetch/test_generic_relationship.py  | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 tests/batch_fetch/test_generic_relationship.py

diff --git a/tests/batch_fetch/test_generic_relationship.py b/tests/batch_fetch/test_generic_relationship.py
new file mode 100644
index 0000000..0748c2c
--- /dev/null
+++ b/tests/batch_fetch/test_generic_relationship.py
@@ -0,0 +1,44 @@
+from __future__ import unicode_literals
+import sqlalchemy as sa
+from tests import TestCase
+from sqlalchemy_utils import batch_fetch, generic_relationship
+
+
+class TestBatchFetchGenericRelationship(TestCase):
+    def create_models(self):
+        class Building(self.Base):
+            __tablename__ = 'building'
+            id = sa.Column(sa.Integer, primary_key=True)
+
+        class User(self.Base):
+            __tablename__ = 'user'
+            id = sa.Column(sa.Integer, primary_key=True)
+
+        class Event(self.Base):
+            __tablename__ = 'event'
+            id = sa.Column(sa.Integer, primary_key=True)
+
+            object_type = sa.Column(sa.Unicode(255))
+            object_id = sa.Column(sa.Integer, nullable=False)
+
+            object = generic_relationship(object_type, object_id)
+
+        self.Building = Building
+        self.User = User
+        self.Event = Event
+
+    def test_batch_fetch(self):
+        user = self.User()
+
+        self.session.add(user)
+        self.session.commit()
+        event = self.Event(object=user)
+        self.session.add(event)
+        self.session.commit()
+
+        events = self.session.query(self.Event).all()
+        batch_fetch(events, 'object')
+        query_count = self.connection.query_count
+        events[0].object
+
+        assert self.connection.query_count == query_count