From 1ee044caa9d7d5509c96e710f18cbcaa361bc45f Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Sat, 24 Aug 2013 13:45:07 +0300
Subject: [PATCH] Added failing composite keys test

---
 sqlalchemy_utils/functions/batch_fetch.py |  3 +-
 tests/batch_fetch/test_composite_keys.py  | 73 +++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 tests/batch_fetch/test_composite_keys.py

diff --git a/sqlalchemy_utils/functions/batch_fetch.py b/sqlalchemy_utils/functions/batch_fetch.py
index 19796c0..fc73282 100644
--- a/sqlalchemy_utils/functions/batch_fetch.py
+++ b/sqlalchemy_utils/functions/batch_fetch.py
@@ -11,7 +11,8 @@ from sqlalchemy.orm.session import object_session
 class with_backrefs(object):
     """
     Marks given attribute path so that whenever its fetched with batch_fetch
-    the backref relations are force set too.
+    the backref relations are force set too. Very useful when dealing with
+    certain many-to-many relationship scenarios.
     """
     def __init__(self, path):
         self.path = path
diff --git a/tests/batch_fetch/test_composite_keys.py b/tests/batch_fetch/test_composite_keys.py
new file mode 100644
index 0000000..d59b80d
--- /dev/null
+++ b/tests/batch_fetch/test_composite_keys.py
@@ -0,0 +1,73 @@
+import sqlalchemy as sa
+from sqlalchemy_utils import batch_fetch
+from tests import TestCase
+
+
+class TestBatchFetchManyToOneRelationships(TestCase):
+    def create_models(self):
+        class User(self.Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(self.Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            name = sa.Column(sa.Unicode(255))
+            author_first_name = sa.Column(
+                sa.Unicode(255), sa.ForeignKey(User.first_name)
+            )
+            author_last_name = sa.Column(
+                sa.Unicode(255), sa.ForeignKey(User.last_name)
+            )
+
+            author = sa.orm.relationship(
+                User,
+                primaryjoin=sa.and_(
+                    author_first_name == User.first_name,
+                    author_last_name == User.last_name
+                ),
+                backref=sa.orm.backref(
+                    'articles'
+                )
+            )
+
+        self.User = User
+        self.Article = Article
+
+    def setup_method(self, method):
+        TestCase.setup_method(self, method)
+        self.users = [
+            self.User(first_name=u'John', last_name=u'Matrix'),
+            self.User(first_name=u'John', last_name=u'The Ripper')
+        ]
+        articles = [
+            self.Article(
+                id=1,
+                name=u'Article 1',
+                author=self.users[0]
+            ),
+            self.Article(
+                id=2,
+                name=u'Article 2',
+                author=self.users[1]
+            ),
+            self.Article(
+                id=3,
+                name=u'Article 3'
+            )
+        ]
+        self.session.add_all(articles)
+        self.session.commit()
+
+    def test_supports_relationship_attributes(self):
+        articles = self.session.query(self.Article).all()
+        batch_fetch(
+            articles,
+            'author'
+        )
+        query_count = self.connection.query_count
+        assert articles[0].author == self.users[0]  # no lazy load should occur
+        assert articles[1].author == self.users[1]  # no lazy load should occur
+        assert articles[2].author is None  # no lazy load should occur
+        assert self.connection.query_count == query_count