diff --git a/asyncio_mongo/connection.py b/asyncio_mongo/connection.py index 1ed02d5..68af4a3 100644 --- a/asyncio_mongo/connection.py +++ b/asyncio_mongo/connection.py @@ -31,13 +31,6 @@ class Connection: connection.host = host - # strip any additional hosts in the connection string for now - try: - if "," in port: - port = port[:port.find(',')] - except TypeError: - pass - connection.port = port connection._loop = loop connection._retry_interval = .5 diff --git a/asyncio_mongo/pool.py b/asyncio_mongo/pool.py index 6a7915f..425fc86 100644 --- a/asyncio_mongo/pool.py +++ b/asyncio_mongo/pool.py @@ -55,7 +55,19 @@ class Pool: raise Exception("Missing database name in URI") self._host = url.hostname - self._port = url.port or 27017 + + try: + port = url.port + except ValueError: + port = url._hostinfo[1] + # strip any additional hosts in the connection string for now + try: + if "," in port: + port = port[:port.find(',')] + except TypeError: + pass + + self._port = port or 27017 username = url.username password = url.password else: diff --git a/tests/test_connection.py b/tests/test_connection.py index 8cd07c6..d479f74 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -43,9 +43,20 @@ class TestMongoConnectionMethods(MongoTest): rapi.close() @async - def test_pool(self): + def test_pool_multiple_hosts(self): # MongoConnectionPool returns deferred, which gets MongoAPI - pool = asyncio_mongo.Pool.create(mongo_host, "%s,blah:333" % mongo_port, poolsize=2) + pool = asyncio_mongo.Pool.create(url="mongodb://{host}:{port},otherhost:333".format( + host=mongo_host, port=mongo_port), poolsize=2) + self.assertTrue(inspect.isgenerator(pool)) + rapi = yield from pool + self.assertEqual(isinstance(rapi, asyncio_mongo.Pool), True) + rapi.close() + + @async + def test_pool_with_url(self): + # MongoConnectionPool returns deferred, which gets MongoAPI + pool = asyncio_mongo.Pool.create(url="mongodb://{host}:{port}".format( + host=mongo_host, port=mongo_port), poolsize=2) self.assertTrue(inspect.isgenerator(pool)) rapi = yield from pool self.assertEqual(isinstance(rapi, asyncio_mongo.Pool), True)