populate default lbp on connection

mambocab · mambocab · commit 4effe4b73b89 · 2017-08-11T11:42:08.000-04:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -4,6 +4,7 @@
 Bug Fixes
 ---------
 * Both _set_final_exception/result called for the same ResponseFuture (PYTHON-630)
+* Use of DCAwareRoundRobinPolicy raises NoHostAvailable exception (PYTHON-781)
 
 
 3.11.0
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -1175,6 +1175,9 @@ def connect(self, keyspace=None, wait_for_all_pools=False):
 
                 self.profile_manager.populate(
                     weakref.proxy(self), self.metadata.all_hosts())
+                self.load_balancing_policy.populate(
+                    weakref.proxy(self), self.metadata.all_hosts()
+                )
 
                 try:
                     self.control_connection.connect()
@@ -2661,7 +2664,13 @@ def _reconnect_internal(self):
         a connection to that host.
         """
         errors = {}
-        for host in self._cluster._default_load_balancing_policy.make_query_plan():
+        lbp = (
+            self._cluster.load_balancing_policy
+            if self._cluster._config_mode == _ConfigMode.LEGACY else
+            self._cluster._default_load_balancing_policy
+        )
+
+        for host in lbp.make_query_plan():
             try:
                 return self._try_connect(host)
             except ConnectionException as exc:
diff --git a/tests/integration/standard/test_cluster.py b/tests/integration/standard/test_cluster.py
@@ -824,6 +824,18 @@ def test_profile_load_balancing(self):
             # make sure original profile is not impacted
             self.assertTrue(session.execute(query, execution_profile='node1')[0].release_version)
 
+    def test_setting_lbp_legacy(self):
+        cluster = Cluster()
+        self.addCleanup(cluster.shutdown)
+        cluster.load_balancing_policy = RoundRobinPolicy()
+        self.assertEqual(
+            list(cluster.load_balancing_policy.make_query_plan()), []
+        )
+        cluster.connect()
+        self.assertNotEqual(
+            list(cluster.load_balancing_policy.make_query_plan()), []
+        )
+
     def test_profile_lb_swap(self):
         """
         Tests that profile load balancing policies are not shared
