Python test_connection.get_connection函数代码示例

    def test_multiple_connections(self):
        a = get_connection(auto_start_request=False)
        b = get_connection(auto_start_request=False)
        self.assertEqual(1, len(a._Connection__pool.sockets))
        self.assertEqual(1, len(b._Connection__pool.sockets))

        a.start_request()
        a.test.test.find_one()
        self.assertEqual(0, len(a._Connection__pool.sockets))
        a.end_request()
        self.assertEqual(1, len(a._Connection__pool.sockets))
        self.assertEqual(1, len(b._Connection__pool.sockets))
        a_sock = one(a._Connection__pool.sockets)

        b.end_request()
        self.assertEqual(1, len(a._Connection__pool.sockets))
        self.assertEqual(1, len(b._Connection__pool.sockets))

        b.start_request()
        b.test.test.find_one()
        self.assertEqual(1, len(a._Connection__pool.sockets))
        self.assertEqual(0, len(b._Connection__pool.sockets))

        b.end_request()
        b_sock = one(b._Connection__pool.sockets)
        b.test.test.find_one()
        a.test.test.find_one()

        self.assertEqual(b_sock,
                         b._Connection__pool.get_socket((b.host, b.port)))
        self.assertEqual(a_sock,
                         a._Connection__pool.get_socket((a.host, a.port)))

    def test_mongos_connection(self):
        c = get_connection()
        is_mongos = utils.is_mongos(c)

        # Test default mode, PRIMARY
        cursor = c.pymongo_test.test.find()
        if is_mongos:
            self.assertEqual(
                {'mode': 'primary'},
                cursor._Cursor__query_spec().get('$readPreference')
            )
        else:
            self.assertFalse(
                '$readPreference' in cursor._Cursor__query_spec())

        # Test non-PRIMARY modes which can be combined with tags
        for mode, mongos_mode in (
            (ReadPreference.PRIMARY_PREFERRED, 'primaryPreferred'),
            (ReadPreference.SECONDARY, 'secondary'),
            (ReadPreference.SECONDARY_PREFERRED, 'secondaryPreferred'),
            (ReadPreference.NEAREST, 'nearest'),
        ):
            for tag_sets in (
                None, [{}]
            ):
                c = get_connection(
                    read_preference=mode,
                    tag_sets=tag_sets)

                self.assertEqual(is_mongos, c.is_mongos)
                cursor = c.pymongo_test.test.find()
                if is_mongos:
                    self.assertEqual(
                        {'mode': mongos_mode},
                        cursor._Cursor__query_spec().get('$readPreference')
                    )
                else:
                    self.assertFalse(
                        '$readPreference' in cursor._Cursor__query_spec())

            for tag_sets in (
                [{'dc': 'la'}],
                [{'dc': 'la'}, {'dc': 'sf'}],
                [{'dc': 'la'}, {'dc': 'sf'}, {}],
            ):
                c = get_connection(
                    read_preference=mode,
                    tag_sets=tag_sets)

                self.assertEqual(is_mongos, c.is_mongos)
                cursor = c.pymongo_test.test.find()
                if is_mongos:
                    self.assertEqual(
                        {'mode': mongos_mode, 'tags': tag_sets},
                        cursor._Cursor__query_spec().get('$readPreference'))
                else:
                    self.assertFalse(
                        '$readPreference' in cursor._Cursor__query_spec())

    def test_only_secondary_ok_commands_have_read_prefs(self):
        c = get_connection(read_preference=ReadPreference.SECONDARY)
        is_mongos = utils.is_mongos(c)
        if not is_mongos:
            raise SkipTest("Only mongos have read_prefs added to the spec")

        # Ensure secondary_ok_commands have readPreference
        for cmd in secondary_ok_commands:
            if cmd == "mapreduce":  # map reduce is a special case
                continue
            command = SON([(cmd, 1)])
            cursor = c.pymongo_test["$cmd"].find(command.copy())
            # White-listed commands also have to be wrapped in $query
            command = SON([("$query", command)])
            command["$readPreference"] = {"mode": "secondary"}
            self.assertEqual(command, cursor._Cursor__query_spec())

        # map_reduce inline should have read prefs
        command = SON([("mapreduce", "test"), ("out", {"inline": 1})])
        cursor = c.pymongo_test["$cmd"].find(command.copy())
        # White-listed commands also have to be wrapped in $query
        command = SON([("$query", command)])
        command["$readPreference"] = {"mode": "secondary"}
        self.assertEqual(command, cursor._Cursor__query_spec())

        # map_reduce that outputs to a collection shouldn't have read prefs
        command = SON([("mapreduce", "test"), ("out", {"mrtest": 1})])
        cursor = c.pymongo_test["$cmd"].find(command.copy())
        self.assertEqual(command, cursor._Cursor__query_spec())

        # Other commands shouldn't be changed
        for cmd in ("drop", "create", "any-future-cmd"):
            command = SON([(cmd, 1)])
            cursor = c.pymongo_test["$cmd"].find(command.copy())
            self.assertEqual(command, cursor._Cursor__query_spec())

    def test_authenticate_and_request(self):
        if (is_mongos(self.connection) and not
            version.at_least(self.connection, (2, 0, 0))):
            raise SkipTest("Auth with sharding requires MongoDB >= 2.0.0")
        # Database.authenticate() needs to be in a request - check that it
        # always runs in a request, and that it restores the request state
        # (in or not in a request) properly when it's finished.
        self.assertTrue(self.connection.auto_start_request)
        db = self.connection.pymongo_test
        db.system.users.remove({})
        db.remove_user("mike")
        db.add_user("mike", "password")
        self.assertTrue(self.connection.in_request())
        self.assertTrue(db.authenticate("mike", "password"))
        self.assertTrue(self.connection.in_request())

        no_request_cx = get_connection(auto_start_request=False)
        no_request_db = no_request_cx.pymongo_test
        self.assertFalse(no_request_cx.in_request())
        self.assertTrue(no_request_db.authenticate("mike", "password"))
        self.assertFalse(no_request_cx.in_request())

        # just make sure there are no exceptions here
        db.logout()
        no_request_db.logout()

    def test_only_secondary_ok_commands_have_read_prefs(self):
        c = get_connection(read_preference=ReadPreference.SECONDARY)
        is_mongos = utils.is_mongos(c)
        if not is_mongos:
            raise SkipTest("Only mongos have read_prefs added to the spec")

        # Ensure secondary_ok_commands have readPreference
        for cmd in secondary_ok_commands:
            if cmd == 'mapreduce':  # map reduce is a special case
                continue
            command = SON([(cmd, 1)])
            cursor = c.pymongo_test["$cmd"].find(command.copy())
            command['$readPreference'] = {'mode': 'secondary'}
            self.assertEqual(command, cursor._Cursor__query_spec())

        # map_reduce inline should have read prefs
        command = SON([('mapreduce', 'test'), ('out', {'inline': 1})])
        cursor = c.pymongo_test["$cmd"].find(command.copy())
        command['$readPreference'] = {'mode': 'secondary'}
        self.assertEqual(command, cursor._Cursor__query_spec())

        # map_reduce that outputs to a collection shouldn't have read prefs
        command = SON([('mapreduce', 'test'), ('out', {'mrtest': 1})])
        cursor = c.pymongo_test["$cmd"].find(command.copy())
        self.assertEqual(command, cursor._Cursor__query_spec())

        # Other commands shouldn't be changed
        for cmd in ('drop', 'create', 'any-future-cmd'):
            command = SON([(cmd, 1)])
            cursor = c.pymongo_test["$cmd"].find(command.copy())
            self.assertEqual(command, cursor._Cursor__query_spec())

    def test_uuid_queries(self):
        if not should_test_uuid:
            raise SkipTest()

        c = get_connection()
        coll = c.pymongo_test.test
        coll.drop()

        uu = uuid.uuid4()
        # Wrap uu.bytes in binary_type to work
        # around http://bugs.python.org/issue7380.
        coll.insert({'uuid': Binary(binary_type(uu.bytes), 3)})
        self.assertEqual(1, coll.count())

        # Test UUIDLegacy queries.
        coll.uuid_subtype = 4
        self.assertEqual(0, coll.find({'uuid': uu}).count())
        cur = coll.find({'uuid': UUIDLegacy(uu)})
        self.assertEqual(1, cur.count())
        retrieved = cur.next()
        self.assertEqual(uu, retrieved['uuid'])

        # Test regular UUID queries (using subtype 4).
        coll.insert({'uuid': uu})
        self.assertEqual(2, coll.count())
        cur = coll.find({'uuid': uu})
        self.assertEqual(1, cur.count())
        retrieved = cur.next()
        self.assertEqual(uu, retrieved['uuid'])

        # Test both.
        cur = coll.find({'uuid': {'$in': [uu, UUIDLegacy(uu)]}})
        self.assertEqual(2, cur.count())
        coll.drop()

 def _get_connection(self):
     """
     Intended for overriding in TestThreadsAuthReplicaSet. This method
     returns a Connection here, and a ReplicaSetConnection in
     test_threads_replica_set_connection.py.
     """
     # Regular test connection
     return get_connection()

 def setUp(self):
     self.db = get_connection().pymongo_test
     self.db.drop_collection("fs.files")
     self.db.drop_collection("fs.chunks")
     self.db.drop_collection("alt.files")
     self.db.drop_collection("alt.chunks")
     self.fs = gridfs.GridFS(self.db)
     self.alt = gridfs.GridFS(self.db, "alt")

    def setUp(self):
        self.c = get_connection(auto_start_request=False)

        # reset the db
        db = self.c[DB]
        db.unique.drop()
        db.test.drop()
        db.unique.insert({"_id": "mike"})
        db.test.insert([{} for i in range(1000)])

def teardown():
    c = get_connection()

    c.drop_database("pymongo-pooling-tests")
    c.drop_database("pymongo_test")
    c.drop_database("pymongo_test1")
    c.drop_database("pymongo_test2")
    c.drop_database("pymongo_test_mike")
    c.drop_database("pymongo_test_bernie")

    def test_authenticate_multiple(self):
        conn = get_connection()
        if (is_mongos(conn) and not
            version.at_least(self.connection, (2, 0, 0))):
            raise SkipTest("Auth with sharding requires MongoDB >= 2.0.0")
        if not server_started_with_auth(conn):
            raise SkipTest("Authentication is not enabled on server")

        # Setup
        users_db = conn.pymongo_test
        admin_db = conn.admin
        other_db = conn.pymongo_test1
        users_db.system.users.remove(safe=True)
        admin_db.system.users.remove(safe=True)
        users_db.test.remove(safe=True)
        other_db.test.remove(safe=True)

        admin_db.add_user('admin', 'pass')
        self.assertTrue(admin_db.authenticate('admin', 'pass'))

        admin_db.add_user('ro-admin', 'pass', read_only=True)
        users_db.add_user('user', 'pass')

        admin_db.logout()
        self.assertRaises(OperationFailure, users_db.test.find_one)

        # Regular user should be able to query its own db, but
        # no other.
        users_db.authenticate('user', 'pass')
        self.assertEqual(0, users_db.test.count())
        self.assertRaises(OperationFailure, other_db.test.find_one)

        # Admin read-only user should be able to query any db,
        # but not write.
        admin_db.authenticate('ro-admin', 'pass')
        self.assertEqual(0, other_db.test.count())
        self.assertRaises(OperationFailure,
                          other_db.test.insert, {}, safe=True)

        # Force close all sockets
        conn.disconnect()

        # We should still be able to write to the regular user's db
        self.assertTrue(users_db.test.remove(safe=True))
        # And read from other dbs...
        self.assertEqual(0, other_db.test.count())
        # But still not write to other dbs...
        self.assertRaises(OperationFailure,
                          other_db.test.insert, {}, safe=True)

        # Cleanup
        admin_db.logout()
        users_db.logout()
        self.assertTrue(admin_db.authenticate('admin', 'pass'))
        self.assertTrue(admin_db.system.users.remove(safe=True))
        self.assertEqual(0, admin_db.system.users.count())
        self.assertTrue(users_db.system.users.remove(safe=True))

    def test_pool_with_fork(self):
        # Test that separate Connections have separate Pools, and that the
        # driver can create a new Connection after forking
        if sys.platform == "win32":
            raise SkipTest("Can't test forking on Windows")

        try:
            from multiprocessing import Process, Pipe
        except ImportError:
            raise SkipTest("No multiprocessing module")

        a = get_connection(auto_start_request=False)
        a.test.test.remove(safe=True)
        a.test.test.insert({'_id':1}, safe=True)
        a.test.test.find_one()
        self.assertEqual(1, len(a._Connection__pool.sockets))
        a_sock = one(a._Connection__pool.sockets)

        def loop(pipe):
            c = get_connection(auto_start_request=False)
            self.assertEqual(1,len(c._Connection__pool.sockets))
            c.test.test.find_one()
            self.assertEqual(1,len(c._Connection__pool.sockets))
            pipe.send(one(c._Connection__pool.sockets).sock.getsockname())

        cp1, cc1 = Pipe()
        cp2, cc2 = Pipe()

        p1 = Process(target=loop, args=(cc1,))
        p2 = Process(target=loop, args=(cc2,))

        p1.start()
        p2.start()

        p1.join(1)
        p2.join(1)

        p1.terminate()
        p2.terminate()

        p1.join()
        p2.join()

        cc1.close()
        cc2.close()

        b_sock = cp1.recv()
        c_sock = cp2.recv()
        self.assertTrue(a_sock.sock.getsockname() != b_sock)
        self.assertTrue(a_sock.sock.getsockname() != c_sock)
        self.assertTrue(b_sock != c_sock)
        self.assertEqual(a_sock,
                         a._Connection__pool.get_socket((a.host, a.port)))

    def test_file_exists(self):
        db = get_connection(w=1).pymongo_test
        fs = gridfs.GridFS(db)

        oid = fs.put(b("hello"))
        self.assertRaises(FileExists, fs.put, b("world"), _id=oid)

        one = fs.new_file(_id=123)
        one.write(b("some content"))
        one.close()

        two = fs.new_file(_id=123)
        self.assertRaises(FileExists, two.write, b('x' * 262146))

    def test_low_network_timeout(self):
        db = None
        i = 0
        n = 10
        while db is None and i < n:
            try:
                db = get_connection(network_timeout=0.0001).pymongo_test
            except AutoReconnect:
                i += 1
        if i == n:
            raise SkipTest()

        threads = []
        for _ in range(4):
            t = IgnoreAutoReconnect(db.test, 100)
            t.start()
            threads.append(t)

        for t in threads:
            t.join()

    def test_authenticate_and_request(self):
        # Database.authenticate() needs to be in a request - check that it
        # always runs in a request, and that it restores the request state
        # (in or not in a request) properly when it's finished.
        self.assertTrue(self.connection.auto_start_request)
        db = self.connection.pymongo_test
        db.system.users.remove({})
        db.remove_user("mike")
        db.add_user("mike", "password")
        self.assertTrue(self.connection.in_request())
        self.assertTrue(db.authenticate("mike", "password"))
        self.assertTrue(self.connection.in_request())

        no_request_cx = get_connection(auto_start_request=False)
        no_request_db = no_request_cx.pymongo_test
        self.assertFalse(no_request_cx.in_request())
        self.assertTrue(no_request_db.authenticate("mike", "password"))
        self.assertFalse(no_request_cx.in_request())

        # just make sure there are no exceptions here
        db.logout()
        no_request_db.logout()