Factor Max Staleness and Server Selection tests

Author: ajdavis

test/test_max_staleness.py

@@ -14,215 +14,31 @@
 
 """Test maxStalenessSeconds support."""
 
-import datetime
 import os
-import time
 import sys
+import time
 import warnings
 
 sys.path[0:0] = [""]
 
-from bson import json_util
-from pymongo import MongoClient, read_preferences
-from pymongo.common import clean_node, HEARTBEAT_FREQUENCY
-from pymongo.errors import ConfigurationError, ConnectionFailure
-from pymongo.ismaster import IsMaster
-from pymongo.server_description import ServerDescription
+from pymongo import MongoClient
+from pymongo.errors import ConfigurationError
 from pymongo.server_selectors import writable_server_selector
-from pymongo.settings import TopologySettings
-from pymongo.topology import Topology
 
 from test import client_context, unittest
 from test.utils import rs_or_single_client
+from test.utils_selection_tests import create_selection_tests
 
 # Location of JSON test specifications.
 _TEST_PATH = os.path.join(
     os.path.dirname(os.path.realpath(__file__)),
     'max_staleness')
 
 
-class MockSocketInfo(object):
-    def close(self):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-
-class MockPool(object):
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def reset(self):
-        pass
-
-
-class MockMonitor(object):
-    def __init__(self, server_description, topology, pool, topology_settings):
-        pass
-
-    def open(self):
-        pass
-
-    def request_check(self):
-        pass
-
-    def close(self):
-        pass
-
-
-def get_addresses(server_list):
-    seeds = []
-    hosts = []
-    for server in server_list:
-        seeds.append(clean_node(server['address']))
-        hosts.append(server['address'])
-    return seeds, hosts
-
-
-def make_last_write_date(server):
-    epoch = datetime.datetime.utcfromtimestamp(0)
-    millis = server.get('lastWrite', {}).get('lastWriteDate')
-    if millis:
-        diff = ((millis % 1000) + 1000) % 1000
-        seconds = (millis - diff) / 1000
-        micros = diff * 1000
-        return epoch + datetime.timedelta(
-            seconds=seconds, microseconds=micros)
-    else:
-        # "Unknown" server.
-        return epoch
-
-
-def make_server_description(server, hosts):
-    """Make ServerDescription from server info from JSON file."""
-    server_type = server['type']
-    if server_type == "Unknown":
-        return ServerDescription(clean_node(server['address']), IsMaster({}))
-
-    ismaster_response = {'ok': True, 'hosts': hosts}
-    if server_type != "Standalone" and server_type != "Mongos":
-        ismaster_response['setName'] = "rs"
-
-    if server_type == "RSPrimary":
-        ismaster_response['ismaster'] = True
-    elif server_type == "RSSecondary":
-        ismaster_response['secondary'] = True
-    elif server_type == "Mongos":
-        ismaster_response['msg'] = 'isdbgrid'
-
-    ismaster_response['lastWrite'] = {
-        'lastWriteDate': make_last_write_date(server)
-    }
-
-    for field in 'maxWireVersion', 'tags', 'idleWritePeriodMillis':
-        if field in server:
-            ismaster_response[field] = server[field]
-
-    # Sets _last_update_time to now.
-    sd = ServerDescription(clean_node(server['address']),
-                           IsMaster(ismaster_response),
-                           round_trip_time=server['avg_rtt_ms'])
-
-    sd._last_update_time = server['lastUpdateTime'] / 1000.0  # ms to sec.
-    return sd
-
-
-class TestAllScenarios(unittest.TestCase):
+class TestAllScenarios(create_selection_tests(_TEST_PATH)):
     pass
 
 
-def create_test(scenario_def):
-    def run_scenario(self):
-        if 'heartbeatFrequencyMS' in scenario_def:
-            frequency = int(scenario_def['heartbeatFrequencyMS']) / 1000.0
-        else:
-            frequency = HEARTBEAT_FREQUENCY
-
-        # Initialize topologies.
-        seeds, hosts = get_addresses(
-            scenario_def['topology_description']['servers'])
-
-        topology = Topology(
-            TopologySettings(seeds=seeds,
-                             monitor_class=MockMonitor,
-                             pool_class=MockPool,
-                             heartbeat_frequency=frequency))
-
-        # Update topologies with server descriptions.
-        for server in scenario_def['topology_description']['servers']:
-            server_description = make_server_description(server, hosts)
-            topology.on_change(server_description)
-
-        # Create server selector.
-        # Make first letter lowercase to match read_pref's modes.
-        pref_def = scenario_def['read_preference']
-        mode_string = pref_def.get('mode', 'primary')
-        mode_string = mode_string[:1].lower() + mode_string[1:]
-        mode = read_preferences.read_pref_mode_from_name(mode_string)
-        max_staleness = pref_def.get('maxStalenessSeconds', -1)
-        tag_sets = pref_def.get('tag_sets')
-
-        if scenario_def.get('error'):
-            with self.assertRaises(ConfigurationError):
-                # Error can be raised when making Read Pref or selecting.
-                pref = read_preferences.make_read_preference(
-                    mode, tag_sets=tag_sets, max_staleness=max_staleness)
-
-                topology.select_server(pref)
-            return
-
-        expected_addrs = set([
-            server['address'] for server in scenario_def['in_latency_window']])
-
-        # Select servers.
-        pref = read_preferences.make_read_preference(
-            mode, tag_sets=tag_sets, max_staleness=max_staleness)
-
-        if not expected_addrs:
-            with self.assertRaises(ConnectionFailure):
-                topology.select_servers(pref, server_selection_timeout=0)
-            return
-
-        servers = topology.select_servers(pref, server_selection_timeout=0)
-        actual_addrs = set(['%s:%d' % s.description.address for s in servers])
-
-        for unexpected in actual_addrs - expected_addrs:
-            self.fail("'%s' shouldn't have been selected, but was" % unexpected)
-
-        for unselected in expected_addrs - actual_addrs:
-            self.fail("'%s' should have been selected, but wasn't" % unselected)
-
-    return run_scenario
-
-
-def create_tests():
-    for dirpath, _, filenames in os.walk(_TEST_PATH):
-        dirname = os.path.split(dirpath)
-        dirname = os.path.split(dirname[-2])[-1] + '_' + dirname[-1]
-
-        for filename in filenames:
-            if not filename.endswith('.json'):
-                continue
-
-            with open(os.path.join(dirpath, filename)) as scenario_stream:
-                scenario_def = json_util.loads(scenario_stream.read())
-
-            # Construct test from scenario.
-            new_test = create_test(scenario_def)
-            test_name = 'test_%s_%s' % (
-                dirname, os.path.splitext(filename)[0])
-
-            new_test.__name__ = test_name
-            setattr(TestAllScenarios, new_test.__name__, new_test)
-
-
-create_tests()
-
-
 class TestMaxStaleness(unittest.TestCase):
     def test_max_staleness(self):
         client = MongoClient()