Avoid AssertionError on first write attempt during replica set failover, PYTHON-601.

ajdavis · ajdavis · commit baf6661a41d8 · 2013-11-26T16:22:40.000-05:00
diff --git a/pymongo/pool.py b/pymongo/pool.py
@@ -393,8 +393,8 @@ def discard_socket(self, sock_info):
     def maybe_return_socket(self, sock_info):
         """Return the socket to the pool unless it's the request socket.
         """
-        # These sentinel values should only be used internally.
-        assert sock_info not in (NO_REQUEST, NO_SOCKET_YET)
+        if sock_info in (NO_REQUEST, NO_SOCKET_YET):
+            return
 
         if self.pid != os.getpid():
             if not sock_info.forced:
diff --git a/test/high_availability/test_ha.py b/test/high_availability/test_ha.py
@@ -388,6 +388,9 @@ def setUp(self):
         res = ha_tools.start_replica_set([{}, {}, {}])
         self.seed, self.name = res
 
+        # Disable periodic refresh.
+        Monitor._refresh_interval = 1e6
+
     def test_writes_with_failover(self):
         c = MongoReplicaSetClient(
             self.seed, replicaSet=self.name, use_greenlets=use_greenlets)
@@ -398,21 +401,35 @@ def test_writes_with_failover(self):
         db.test.insert({'foo': 'bar'}, w=w)
         self.assertEqual('bar', db.test.find_one()['foo'])
 
-        def try_write():
-            for _ in xrange(30):
-                try:
-                    db.test.insert({'bar': 'baz'})
-                    return True
-                except AutoReconnect:
-                    sleep(1)
-            return False
-
         killed = ha_tools.kill_primary(9)
         self.assertTrue(bool(len(killed)))
-        self.assertTrue(try_write())
+
+        # Wait past pool's check interval, so it throws an error from
+        # get_socket().
+        sleep(1)
+
+        # Verify that we only raise AutoReconnect, not some other error,
+        # while we wait for new primary.
+        for _ in xrange(10000):
+            try:
+                db.test.insert({'bar': 'baz'})
+
+                # No error, found primary.
+                break
+            except AutoReconnect:
+                sleep(.01)
+        else:
+            self.fail("Couldn't connect to new primary")
+
+        # Found new primary.
+        self.assertTrue(c.primary)
         self.assertTrue(primary != c.primary)
         self.assertEqual('baz', db.test.find_one({'bar': 'baz'})['bar'])
 
+    def tearDown(self):
+        Monitor._refresh_interval = MONITOR_INTERVAL
+        super(TestWritesWithFailover, self).tearDown()
+
 
 class TestReadWithFailover(HATestCase):
 
