PYTHON-2462 Implement prose test "Connection Pool Management"

Author: ShaneHarvey

test/test_discovery_and_monitoring.py

@@ -39,11 +39,14 @@
 from test import unittest, IntegrationTest
 from test.utils import (assertion_context,
                         cdecimal_patched,
+                        CMAPListener,
                         client_context,
                         Barrier,
                         get_pool,
+                        HeartbeatEventListener,
                         server_name_to_type,
                         rs_or_single_client,
+                        single_client,
                         TestCreator,
                         wait_until)
 from test.utils_spec_runner import SpecRunner, SpecRunnerThread
@@ -299,6 +302,46 @@ def insert_command(i):
         client.admin.command('ping')
 
 
+class CMAPHeartbeatListener(HeartbeatEventListener, CMAPListener):
+    pass
+
+
+class TestPoolManagement(IntegrationTest):
+    @client_context.require_failCommand_appName
+    def test_pool_unpause(self):
+        # This test implements the prose test "Connection Pool Management"
+        listener = CMAPHeartbeatListener()
+        client = single_client(appName="SDAMPoolManagementTest",
+                               heartbeatFrequencyMS=500,
+                               event_listeners=[listener])
+        self.addCleanup(client.close)
+        # Assert that ConnectionPoolReadyEvent occurs after the first
+        # ServerHeartbeatSucceededEvent.
+        listener.wait_for_event(monitoring.PoolReadyEvent, 1)
+        pool_ready = listener.events_by_type(monitoring.PoolReadyEvent)[0]
+        hb_succeeded = listener.events_by_type(
+            monitoring.ServerHeartbeatSucceededEvent)[0]
+        self.assertGreater(
+            listener.events.index(pool_ready),
+            listener.events.index(hb_succeeded))
+
+        listener.reset()
+        fail_ismaster = {
+            'mode': {'times': 2},
+            'data': {
+                'failCommands': ['isMaster'],
+                'errorCode': 1234,
+                'appName': 'SDAMPoolManagementTest',
+            },
+        }
+        with self.fail_point(fail_ismaster):
+            listener.wait_for_event(monitoring.ServerHeartbeatFailedEvent, 1)
+            listener.wait_for_event(monitoring.PoolClearedEvent, 1)
+            listener.wait_for_event(
+                monitoring.ServerHeartbeatSucceededEvent, 1)
+            listener.wait_for_event(monitoring.PoolReadyEvent, 1)
+
+
 class TestIntegration(SpecRunner):
     # Location of JSON test specifications.
     TEST_PATH = os.path.join(

test/test_heartbeat_monitoring.py

@@ -51,12 +51,12 @@ def _check_with_socket(self, *args, **kwargs):
             # monitor thread may run multiple times during the execution
             # of this test.
             wait_until(
-                lambda: len(listener.results) >= expected_len,
+                lambda: len(listener.events) >= expected_len,
                 "publish all events")
 
         try:
             # zip gives us len(expected_results) pairs.
-            for expected, actual in zip(expected_results, listener.results):
+            for expected, actual in zip(expected_results, listener.events):
                 self.assertEqual(expected,
                                  actual.__class__.__name__)
                 self.assertEqual(actual.connection_id,

test/test_streaming_protocol.py

@@ -203,7 +203,7 @@ def hb_failed(event):
         self.assertTrue(hb_failed_events[0].awaited)
         # Depending on thread scheduling, the failed heartbeat could occur on
         # the second or third check.
-        events = [type(e) for e in hb_listener.results[:4]]
+        events = [type(e) for e in hb_listener.events[:4]]
         if events == [monitoring.ServerHeartbeatStartedEvent,
                       monitoring.ServerHeartbeatSucceededEvent,
                       monitoring.ServerHeartbeatStartedEvent,

test/utils.py

@@ -60,7 +60,7 @@
 IMPOSSIBLE_WRITE_CONCERN = WriteConcern(w=50)
 
 
-class CMAPListener(ConnectionPoolListener):
+class BaseListener(ConnectionPoolListener):
     def __init__(self):
         self.events = []
 
@@ -71,9 +71,26 @@ def add_event(self, event):
         self.events.append(event)
 
     def event_count(self, event_type):
-        return len([event for event in self.events[:]
-                    if isinstance(event, event_type)])
+        return len(self.events_by_type(event_type))
 
+    def events_by_type(self, event_type):
+        """Return the matching events by event class.
+
+        event_type can be a single class or a tuple of classes.
+        """
+        return self.matching(lambda e: isinstance(e, event_type))
+
+    def matching(self, matcher):
+        """Return the matching events."""
+        return [event for event in self.events[:] if matcher(event)]
+
+    def wait_for_event(self, event, count):
+        """Wait for a number of events to be published, or fail."""
+        wait_until(lambda: self.event_count(event) >= count,
+                   'find %s %s event(s)' % (count, event))
+
+
+class CMAPListener(BaseListener, monitoring.ConnectionPoolListener):
     def connection_created(self, event):
         self.add_event(event)
 
@@ -199,25 +216,17 @@ class ServerAndTopologyEventListener(ServerEventListener,
     """Listens to Server and Topology events."""
 
 
-class HeartbeatEventListener(monitoring.ServerHeartbeatListener):
+class HeartbeatEventListener(BaseListener, monitoring.ServerHeartbeatListener):
     """Listens to only server heartbeat events."""
 
-    def __init__(self):
-        self.results = []
-
     def started(self, event):
-        self.results.append(event)
+        self.add_event(event)
 
     def succeeded(self, event):
-        self.results.append(event)
+        self.add_event(event)
 
     def failed(self, event):
-        self.results.append(event)
-
-    def matching(self, matcher):
-        """Return the matching events."""
-        results = self.results[:]
-        return [event for event in results if matcher(event)]
+        self.add_event(event)
 
 
 class MockSocketInfo(object):