# Copyright 2015 MongoDB, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Test PyMongo's SlaveOkay with:

- A direct connection to a standalone.
- A direct connection to a slave.
- A direct connection to a mongos.
"""
import itertools

from pymongo.read_preferences import make_read_preference
from pymongo.read_preferences import read_pref_mode_from_name

try:
    from queue import Queue
except ImportError:
    from Queue import Queue

from mockupdb import MockupDB, going, QUERY_FLAGS
from pymongo import MongoClient

from tests import unittest
from tests.operations import operations


class TestSlaveOkaySharded(unittest.TestCase):
    def setUp(self):
        self.mongos1, self.mongos2 = MockupDB(), MockupDB()

        # Collect queries to either server in one queue.
        self.q = Queue()
        for server in self.mongos1, self.mongos2:
            server.subscribe(self.q.put)
            server.run()
            self.addCleanup(server.stop)
            server.autoresponds('ismaster', ismaster=True, msg='isdbgrid')

        self.mongoses_uri = 'mongodb://%s,%s' % (self.mongos1.address_string,
                                                 self.mongos2.address_string)


def create_slave_ok_sharded_test(mode, operation):
    def test(self):
        if operation.op_type == 'always-use-secondary':
            slave_ok = True
        elif operation.op_type == 'may-use-secondary':
            slave_ok = mode != 'primary'
        elif operation.op_type == 'must-use-primary':
            slave_ok = False
        else:
            assert False, 'unrecognized op_type %r' % operation.op_type

        pref = make_read_preference(read_pref_mode_from_name(mode),
                                    tag_sets=None)

        client = MongoClient(self.mongoses_uri, read_preference=pref)
        with going(operation.function, client):
            request = self.q.get(timeout=1)
            request.reply(operation.reply)

        if slave_ok:
            self.assertTrue(request.flags & QUERY_FLAGS['SlaveOkay'],
                            'SlaveOkay not set')
        else:
            self.assertFalse(request.flags & QUERY_FLAGS['SlaveOkay'],
                             'SlaveOkay set')

    return test


def generate_slave_ok_sharded_tests():
    modes = 'primary', 'secondary', 'nearest'
    matrix = itertools.product(modes, operations)

    for entry in matrix:
        mode, operation = entry
        test = create_slave_ok_sharded_test(mode, operation)
        test_name = 'test_%s_with_mode_%s' % (
            operation.name.replace(' ', '_'), mode)

        test.__name__ = test_name
        setattr(TestSlaveOkaySharded, test_name, test)


generate_slave_ok_sharded_tests()

if __name__ == '__main__':
    unittest.main()

@@ -1149,15 +1149,36 @@
  set
-Up(self
+up_server(self, wire_version
 ):%0A 
@@ -1519,16 +1519,77 @@
 master',
+ maxWireVersion=wire_version,%0A                               
  ismaste
@@ -1818,24 +1818,24 @@
 operation):%0A
-
     def test
@@ -1842,16 +1842,66 @@
 (self):%0A
+        self.setup_server(operation.wire_version)%0A
         

Set server wire version in sharded mode.