From 2dde984e88e1b30144f35e11dc12824bf3e470aa Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Wed, 3 Oct 2012 11:10:42 -0400
Subject: [PATCH] - add one more #2583 test to cover the "multiple PK switch"
 use case

---
 lib/sqlalchemy/orm/session.py |  2 +-
 test/orm/test_transaction.py  | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index faa9e5a83c..1df9d45ca7 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -1288,7 +1288,7 @@ class Session(_SessionClassMethods):
                     # map (see test/orm/test_naturalpks.py ReversePKsTest)
                     self.identity_map.discard(state)
                     if state in self.transaction._key_switches:
-                        orig_key = self.transaction._key_switches[0]
+                        orig_key = self.transaction._key_switches[state][0]
                     else:
                         orig_key = state.key
                     self.transaction._key_switches[state] = (orig_key, instance_key)
diff --git a/test/orm/test_transaction.py b/test/orm/test_transaction.py
index 49fdd28649..9c19d5bdf3 100644
--- a/test/orm/test_transaction.py
+++ b/test/orm/test_transaction.py
@@ -1184,6 +1184,39 @@ class NaturalPKRollbackTest(fixtures.MappedTest):
         assert s.identity_map[(User, ('u1',))] is u1
         assert s.identity_map[(User, ('u2',))] is u2
 
+    def test_multiple_key_replaced_by_update(self):
+        users, User = self.tables.users, self.classes.User
+
+        mapper(User, users)
+
+        u1 = User(name='u1')
+        u2 = User(name='u2')
+        u3 = User(name='u3')
+
+        s = Session()
+        s.add_all([u1, u2, u3])
+        s.commit()
+
+        s.delete(u1)
+        s.delete(u2)
+        s.flush()
+
+        u3.name = 'u1'
+        s.flush()
+
+        u3.name = 'u2'
+        s.flush()
+
+        s.rollback()
+
+        assert u1 in s
+        assert u2 in s
+        assert u3 in s
+
+        assert s.identity_map[(User, ('u1',))] is u1
+        assert s.identity_map[(User, ('u2',))] is u2
+        assert s.identity_map[(User, ('u3',))] is u3
+
     def test_key_replaced_by_oob_insert(self):
         users, User = self.tables.users, self.classes.User
 
-- 
2.47.3