diff --git a/aiogremlin/client.py b/aiogremlin/client.py index 965a8e24fee410e8921bde856f91e368d39f7b97..5e30106fca64d3547f5308fbd923d67b24394743 100644 --- a/aiogremlin/client.py +++ b/aiogremlin/client.py @@ -70,15 +70,16 @@ class GremlinClient: self.timeout = timeout self._pool = pool self._connector = connector - self._factory = factory or GremlinFactory(connector=self._connector) self._conn = connection if pool is not None: + factory = pool._factory self._connected = self._pool._connected elif self._conn and not getattr(connection, "closed", True): self._connected = True else: self._connected = False self._conn = asyncio.async(self._connect(), loop=self._loop) + self._factory = factory or GremlinFactory(connector=self._connector) if verbose: logger.setLevel(INFO) diff --git a/aiogremlin/pool.py b/aiogremlin/pool.py index c6e90b05d28a368f06fb62476827a395bcb30488..328e780c9f69576acb34779258c3f7d23220aeba 100644 --- a/aiogremlin/pool.py +++ b/aiogremlin/pool.py @@ -16,11 +16,11 @@ class WebSocketPool: """ """ self.url = url - self._ws_response_class = (ws_response_class or - GremlinClientWebSocketResponse) + if ws_response_class is None: + ws_response_class = GremlinClientWebSocketResponse self._factory = factory or GremlinFactory( connector=connector, - ws_response_class=self._ws_response_class) + ws_response_class=ws_response_class) self.poolsize = poolsize self.max_retries = max_retries self.timeout = timeout