octavia 12.0.0.0rc2__py3-none-any.whl → 13.0.0.0rc1__py3-none-any.whl

octavia/tests/unit/common/test_clients.py

@@ -13,7 +13,6 @@ from unittest import mock
 
 import cinderclient.v3
 import glanceclient.v2
-import neutronclient.v2_0
 import novaclient.v2
 from oslo_config import cfg
 
@@ -62,44 +61,6 @@ class TestNovaAuth(base.TestCase):
         self.assertIs(bc1, bc2)
 
 
-class TestNeutronAuth(base.TestCase):
-
-    def setUp(self):
-        # Reset the session and client
-        clients.NeutronAuth.neutron_client = None
-        keystone._SESSION = None
-
-        super().setUp()
-
-    @mock.patch('keystoneauth1.session.Session', mock.Mock())
-    def test_get_neutron_client(self):
-        # There should be no existing client
-        self.assertIsNone(
-            clients.NeutronAuth.neutron_client
-        )
-
-        # Mock out the keystone session and get the client
-        keystone._SESSION = mock.MagicMock()
-        bc1 = clients.NeutronAuth.get_neutron_client(
-            region=None, endpoint_type='publicURL')
-
-        # Our returned client should also be the saved client
-        self.assertIsInstance(
-            clients.NeutronAuth.neutron_client,
-            neutronclient.v2_0.client.Client
-        )
-        self.assertIs(
-            clients.NeutronAuth.neutron_client,
-            bc1
-        )
-
-        # Getting the session again should return the same object
-        bc2 = clients.NeutronAuth.get_neutron_client(
-            region="test-region", service_name="neutronEndpoint1",
-            endpoint="test-endpoint", endpoint_type='publicURL', insecure=True)
-        self.assertIs(bc1, bc2)
-
-
 class TestGlanceAuth(base.TestCase):
 
     def setUp(self):

octavia/tests/unit/common/test_keystone.py

@@ -0,0 +1,54 @@
+#  Copyright Red Hat
+#
+#     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.
+from unittest import mock
+from unittest.mock import call
+
+from keystoneauth1 import exceptions as ks_exceptions
+from oslo_config import cfg
+from oslo_config import fixture as oslo_fixture
+
+import octavia.common.keystone as ks
+import octavia.tests.unit.base as base
+
+
+class TestKeystoneSession(base.TestCase):
+
+    @mock.patch("oslo_config.cfg.ConfigOpts.get_location", return_value=None)
+    @mock.patch("octavia.common.keystone.ks_loading"
+                ".load_auth_from_conf_options")
+    @mock.patch("octavia.common.keystone.LOG")
+    def test_get_auth_neutron_override(self, mock_log, mock_load_auth,
+                                       mock_get_location):
+        opt_mock = mock.MagicMock()
+        opt_mock.dest = "foo"
+        conf = oslo_fixture.Config(cfg.CONF)
+        conf.conf.service_auth.cafile = "bar"
+
+        mock_load_auth.side_effect = [
+            ks_exceptions.auth_plugins.MissingRequiredOptions(
+                [opt_mock]),
+            None,
+            None
+        ]
+
+        sess = ks.KeystoneSession("neutron")
+        sess.get_auth()
+
+        mock_load_auth.assert_has_calls([call(cfg.CONF, 'neutron'),
+                                         call(cfg.CONF, 'service_auth'),
+                                         call(cfg.CONF, 'neutron')])
+        mock_log.debug.assert_has_calls(
+            [call("Overriding [%s].%s with '%s'", 'neutron', 'cafile',
+                  'bar')]
+        )

octavia/tests/unit/common/test_validate.py

@@ -16,6 +16,7 @@ from unittest import mock
 from oslo_config import cfg
 from oslo_config import fixture as oslo_fixture
 from oslo_utils import uuidutils
+from wsme import types as wtypes
 
 import octavia.common.constants as constants
 import octavia.common.exceptions as exceptions
@@ -536,3 +537,69 @@ class TestValidations(base.TestCase):
                                                       '2001:db8::/32'))
         self.assertFalse(validate.is_ip_member_of_cidr('::ffff:0:203.0.113.5',
                                                        '2001:db8::/32'))
+
+    def test_check_hsts_options(self):
+        self.assertRaises(
+            exceptions.ValidationException,
+            validate.check_hsts_options,
+            {'hsts_include_subdomains': True,
+             'hsts_preload': wtypes.Unset,
+             'hsts_max_age': wtypes.Unset}
+        )
+        self.assertRaises(
+            exceptions.ValidationException,
+            validate.check_hsts_options,
+            {'hsts_include_subdomains': wtypes.Unset,
+             'hsts_preload': True,
+             'hsts_max_age': wtypes.Unset}
+        )
+        self.assertRaises(
+            exceptions.ValidationException,
+            validate.check_hsts_options,
+            {'protocol': constants.PROTOCOL_UDP,
+             'hsts_include_subdomains': wtypes.Unset,
+             'hsts_preload': wtypes.Unset,
+             'hsts_max_age': 1}
+        )
+        self.assertIsNone(
+            validate.check_hsts_options(
+                {'protocol': constants.PROTOCOL_TERMINATED_HTTPS,
+                 'hsts_include_subdomains': wtypes.Unset,
+                 'hsts_preload': wtypes.Unset,
+                 'hsts_max_age': 1})
+        )
+
+    def test_check_hsts_options_put(self):
+        listener = mock.MagicMock()
+        db_listener = mock.MagicMock()
+        db_listener.protocol = constants.PROTOCOL_TERMINATED_HTTPS
+
+        listener.hsts_max_age = wtypes.Unset
+        db_listener.hsts_max_age = None
+        for obj in (listener, db_listener):
+            obj.hsts_include_subdomains = False
+            obj.hsts_preload = False
+        self.assertIsNone(validate.check_hsts_options_put(
+            listener, db_listener))
+
+        for i in range(2):
+            listener.hsts_include_subdomains = bool(i % 2)
+            listener.hsts_preload = not bool(i % 2)
+            self.assertRaises(
+                exceptions.ValidationException,
+                validate.check_hsts_options_put,
+                listener, db_listener)
+
+        listener.hsts_max_age, db_listener.hsts_max_age = wtypes.Unset, 0
+        self.assertIsNone(validate.check_hsts_options_put(
+            listener, db_listener))
+
+        listener.hsts_max_age, db_listener.hsts_max_age = 3, None
+        self.assertIsNone(validate.check_hsts_options_put(
+            listener, db_listener))
+
+        db_listener.protocol = constants.PROTOCOL_HTTP
+        self.assertRaises(
+            exceptions.ValidationException,
+            validate.check_hsts_options_put,
+            listener, db_listener)

octavia/tests/unit/controller/healthmanager/test_health_manager.py

@@ -43,15 +43,13 @@ class TestHealthManager(base.TestCase):
         super().setUp()
 
     @mock.patch('octavia.db.api.wait_for_connection')
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.failover_amphora')
     @mock.patch('octavia.controller.worker.v2.controller_worker.'
                 'ControllerWorker.failover_amphora')
     @mock.patch('octavia.db.repositories.AmphoraHealthRepository.'
                 'get_stale_amphora')
     @mock.patch('octavia.db.api.get_session')
     def test_health_check_stale_amphora(self, session_mock, get_stale_amp_mock,
-                                        failover_mockv2, failover_mock,
+                                        failover_mock,
                                         db_wait_mock):
         conf = oslo_fixture.Config(cfg.CONF)
         conf.config(group="health_manager", heartbeat_timeout=5)
@@ -87,15 +85,13 @@ class TestHealthManager(base.TestCase):
         self.assertRaises(TestException, hm.health_check)
         self.assertEqual(4, mock_session.rollback.call_count)
 
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.failover_amphora')
     @mock.patch('octavia.controller.worker.v2.controller_worker.'
                 'ControllerWorker.failover_amphora')
     @mock.patch('octavia.db.repositories.AmphoraHealthRepository.'
                 'get_stale_amphora', return_value=None)
     @mock.patch('octavia.db.api.get_session')
     def test_health_check_nonstale_amphora(self, session_mock,
-                                           get_stale_amp_mock, failover_mockv2,
+                                           get_stale_amp_mock,
                                            failover_mock):
         get_stale_amp_mock.side_effect = [None, TestException('test')]
 
@@ -103,42 +99,32 @@ class TestHealthManager(base.TestCase):
         hm = healthmanager.HealthManager(exit_event)
 
         hm.health_check()
-        session_mock.assert_called_once_with(autocommit=False)
-        if CONF.api_settings.default_provider_driver == 'amphorav2':
-            self.assertFalse(failover_mockv2.called)
-        else:
-            self.assertFalse(failover_mock.called)
+        session_mock.assert_called_once_with()
+        self.assertFalse(failover_mock.called)
 
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.failover_amphora')
     @mock.patch('octavia.controller.worker.v2.controller_worker.'
                 'ControllerWorker.failover_amphora')
     @mock.patch('octavia.db.repositories.AmphoraHealthRepository.'
                 'get_stale_amphora', return_value=None)
     @mock.patch('octavia.db.api.get_session')
     def test_health_check_exit(self, session_mock, get_stale_amp_mock,
-                               failover_mockv2, failover_mock):
+                               failover_mock):
         get_stale_amp_mock.return_value = None
 
         exit_event = threading.Event()
         hm = healthmanager.HealthManager(exit_event)
         hm.health_check()
 
-        session_mock.assert_called_once_with(autocommit=False)
-        if CONF.api_settings.default_provider_driver == 'amphorav2':
-            self.assertFalse(failover_mockv2.called)
-        else:
-            self.assertFalse(failover_mock.called)
+        session_mock.assert_called_once_with()
+        self.assertFalse(failover_mock.called)
 
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.failover_amphora')
     @mock.patch('octavia.controller.worker.v2.controller_worker.'
                 'ControllerWorker.failover_amphora')
     @mock.patch('octavia.db.repositories.AmphoraHealthRepository.'
                 'get_stale_amphora', return_value=None)
     @mock.patch('octavia.db.api.get_session')
     def test_health_check_db_error(self, session_mock, get_stale_amp_mock,
-                                   failover_mockv2, failover_mock):
+                                   failover_mock):
         get_stale_amp_mock.return_value = None
 
         mock_session = mock.MagicMock()

octavia/tests/unit/controller/housekeeping/test_house_keeping.py

@@ -158,73 +158,12 @@ class TestCertRotation(base.TestCase):
         super().setUp()
         self.CONF = self.useFixture(oslo_fixture.Config(cfg.CONF))
 
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.amphora_cert_rotation')
-    @mock.patch('octavia.db.repositories.AmphoraRepository.'
-                'get_cert_expiring_amphora')
-    @mock.patch('octavia.db.api.get_session')
-    def test_cert_rotation_expired_amphora_with_exception(self, session,
-                                                          cert_exp_amp_mock,
-                                                          amp_cert_mock
-                                                          ):
-        self.CONF.config(group="api_settings",
-                         default_provider_driver='amphorav1')
-        amphora = mock.MagicMock()
-        amphora.id = AMPHORA_ID
-
-        session.return_value = session
-        cert_exp_amp_mock.side_effect = [amphora, TestException(
-            'break_while')]
-
-        cr = house_keeping.CertRotation()
-        self.assertRaises(TestException, cr.rotate)
-        amp_cert_mock.assert_called_once_with(AMPHORA_ID)
-
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.amphora_cert_rotation')
-    @mock.patch('octavia.db.repositories.AmphoraRepository.'
-                'get_cert_expiring_amphora')
-    @mock.patch('octavia.db.api.get_session')
-    def test_cert_rotation_expired_amphora_without_exception(self, session,
-                                                             cert_exp_amp_mock,
-                                                             amp_cert_mock
-                                                             ):
-        self.CONF.config(group="api_settings",
-                         default_provider_driver='amphorav1')
-        amphora = mock.MagicMock()
-        amphora.id = AMPHORA_ID
-
-        session.return_value = session
-        cert_exp_amp_mock.side_effect = [amphora, None]
-
-        cr = house_keeping.CertRotation()
-
-        self.assertIsNone(cr.rotate())
-        amp_cert_mock.assert_called_once_with(AMPHORA_ID)
-
-    @mock.patch('octavia.controller.worker.v1.controller_worker.'
-                'ControllerWorker.amphora_cert_rotation')
-    @mock.patch('octavia.db.repositories.AmphoraRepository.'
-                'get_cert_expiring_amphora')
-    @mock.patch('octavia.db.api.get_session')
-    def test_cert_rotation_non_expired_amphora(self, session,
-                                               cert_exp_amp_mock,
-                                               amp_cert_mock):
-        self.CONF.config(group="api_settings",
-                         default_provider_driver='amphorav1')
-
-        session.return_value = session
-        cert_exp_amp_mock.return_value = None
-        cr = house_keeping.CertRotation()
-        cr.rotate()
-        self.assertFalse(amp_cert_mock.called)
-
     @mock.patch('octavia.controller.worker.v2.controller_worker.'
                 'ControllerWorker.amphora_cert_rotation')
     @mock.patch('octavia.db.repositories.AmphoraRepository.'
                 'get_cert_expiring_amphora')
     @mock.patch('octavia.db.api.get_session')
-    def test_cert_rotation_expired_amphora_with_exception_amphorav2(
+    def test_cert_rotation_expired_amphora_with_exception(
             self, session, cert_exp_amp_mock, amp_cert_mock):
         self.CONF.config(group="api_settings",
                          default_provider_driver='amphora')
@@ -245,7 +184,7 @@ class TestCertRotation(base.TestCase):
     @mock.patch('octavia.db.repositories.AmphoraRepository.'
                 'get_cert_expiring_amphora')
     @mock.patch('octavia.db.api.get_session')
-    def test_cert_rotation_expired_amphora_without_exception_amphorav2(
+    def test_cert_rotation_expired_amphora_without_exception(
             self, session, cert_exp_amp_mock, amp_cert_mock):
         self.CONF.config(group="api_settings",
                          default_provider_driver='amphora')
@@ -265,7 +204,7 @@ class TestCertRotation(base.TestCase):
     @mock.patch('octavia.db.repositories.AmphoraRepository.'
                 'get_cert_expiring_amphora')
     @mock.patch('octavia.db.api.get_session')
-    def test_cert_rotation_non_expired_amphora_amphorav2(
+    def test_cert_rotation_non_expired_amphora(
             self, session, cert_exp_amp_mock, amp_cert_mock):
         self.CONF.config(group="api_settings",
                          default_provider_driver='amphora')

octavia/tests/unit/controller/worker/test_amphora_rate_limit.py

@@ -36,7 +36,7 @@ class TestAmphoraBuildRateLimit(base.TestCase):
         self.amp_build_req_repo = mock.MagicMock()
         self.conf.config(group='haproxy_amphora', build_rate_limit=1)
 
-    @mock.patch('octavia.db.api.get_session', mock.MagicMock())
+    @mock.patch('octavia.db.api.session', mock.MagicMock())
     @mock.patch('octavia.controller.worker.amphora_rate_limit'
                 '.AmphoraBuildRateLimit.wait_for_build_slot')
     @mock.patch('octavia.db.repositories.AmphoraBuildReqRepository'

octavia/tests/unit/controller/worker/test_task_utils.py

@@ -19,8 +19,6 @@ from octavia.common import constants
 from octavia.controller.worker import task_utils as task_utilities
 import octavia.tests.unit.base as base
 
-TEST_SESSION = 'TEST_SESSION'
-
 
 class TestTaskUtils(base.TestCase):
 
@@ -39,7 +37,7 @@ class TestTaskUtils(base.TestCase):
 
         super().setUp()
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.AmphoraRepository.update')
     def test_mark_amphora_status_error(self,
                                        mock_amphora_repo_update,
@@ -48,8 +46,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_amphora_status_error(self.AMPHORA_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_amphora_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.AMPHORA_ID,
             status=constants.ERROR)
 
@@ -61,7 +61,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_amphora_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.HealthMonitorRepository.update')
     def test_mark_health_mon_prov_status_error(self,
                                                mock_health_mon_repo_update,
@@ -70,8 +70,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_health_mon_prov_status_error(self.HEALTH_MON_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_health_mon_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.HEALTH_MON_ID,
             provisioning_status=constants.ERROR)
 
@@ -83,7 +85,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_health_mon_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.L7PolicyRepository.update')
     def test_mark_l7policy_prov_status_error(self,
                                              mock_l7policy_repo_update,
@@ -92,8 +94,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_l7policy_prov_status_error(self.L7POLICY_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_l7policy_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.L7POLICY_ID,
             provisioning_status=constants.ERROR)
 
@@ -105,7 +109,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_l7policy_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.L7RuleRepository.update')
     def test_mark_l7rule_prov_status_error(self,
                                            mock_l7rule_repo_update,
@@ -114,8 +118,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_l7rule_prov_status_error(self.L7RULE_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_l7rule_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.L7RULE_ID,
             provisioning_status=constants.ERROR)
 
@@ -127,7 +133,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_l7rule_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.ListenerRepository.update')
     def test_mark_listener_prov_status_active(self,
                                               mock_listener_repo_update,
@@ -136,8 +142,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_listener_prov_status_active(self.LISTENER_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_listener_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.LISTENER_ID,
             provisioning_status=constants.ACTIVE)
 
@@ -149,7 +157,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_listener_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.ListenerRepository.update')
     def test_mark_listener_prov_status_error(self,
                                              mock_listener_repo_update,
@@ -158,8 +166,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_listener_prov_status_error(self.LISTENER_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_listener_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.LISTENER_ID,
             provisioning_status=constants.ERROR)
 
@@ -171,7 +181,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_listener_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.LoadBalancerRepository.update')
     def test_mark_loadbalancer_prov_status_active(self,
                                                   mock_lb_repo_update,
@@ -181,8 +191,10 @@ class TestTaskUtils(base.TestCase):
         self.task_utils.mark_loadbalancer_prov_status_active(
             self.LOADBALANCER_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_lb_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.LOADBALANCER_ID,
             provisioning_status=constants.ACTIVE)
 
@@ -195,7 +207,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_lb_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.LoadBalancerRepository.update')
     def test_mark_loadbalancer_prov_status_error(self,
                                                  mock_lb_repo_update,
@@ -205,8 +217,10 @@ class TestTaskUtils(base.TestCase):
         self.task_utils.mark_loadbalancer_prov_status_error(
             self.LOADBALANCER_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_lb_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.LOADBALANCER_ID,
             provisioning_status=constants.ERROR)
 
@@ -219,7 +233,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_lb_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.MemberRepository.update')
     def test_mark_member_prov_status_error(self,
                                            mock_member_repo_update,
@@ -228,8 +242,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_member_prov_status_error(self.MEMBER_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_member_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.MEMBER_ID,
             provisioning_status=constants.ERROR)
 
@@ -241,7 +257,7 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_member_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.PoolRepository.update')
     def test_mark_pool_prov_status_error(self,
                                          mock_pool_repo_update,
@@ -250,8 +266,10 @@ class TestTaskUtils(base.TestCase):
         # Happy path
         self.task_utils.mark_pool_prov_status_error(self.POOL_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_pool_repo_update.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.POOL_ID,
             provisioning_status=constants.ERROR)
 
@@ -263,15 +281,17 @@ class TestTaskUtils(base.TestCase):
 
         self.assertFalse(mock_pool_repo_update.called)
 
-    @mock.patch('octavia.db.api.get_session', return_value=TEST_SESSION)
+    @mock.patch('octavia.db.api.session')
     @mock.patch('octavia.db.repositories.LoadBalancerRepository.get')
     def test_get_current_loadbalancer_from_db(self, mock_lb_repo_get,
                                               mock_get_session):
         # Happy path
         self.task_utils.get_current_loadbalancer_from_db(self.LOADBALANCER_ID)
 
+        mock_session = mock_get_session().begin().__enter__()
+
         mock_lb_repo_get.assert_called_once_with(
-            TEST_SESSION,
+            mock_session,
             id=self.LOADBALANCER_ID)
 
         # Exception path

octavia/tests/unit/controller/worker/v2/flows/test_load_balancer_flows.py

@@ -208,7 +208,6 @@ class TestLoadBalancerFlows(base.TestCase):
         self.assertEqual(2, len(amp_flow.requires), amp_flow.requires)
         self.assertEqual(4, len(amp_flow.provides), amp_flow.provides)
 
-        # Test mark_active=False
         amp_flow = self.LBFlow.get_post_lb_amp_association_flow(
             '123', constants.TOPOLOGY_ACTIVE_STANDBY)