From 55de6c47f68b945d0b4c8ec294c71b6183fc1b75 Mon Sep 17 00:00:00 2001 From: Yihan Zhao Date: Wed, 20 Nov 2024 16:01:29 +1100 Subject: [PATCH] Fix the index management tests (#1044) --- .../index_management/test_index_management.py | 62 ++++++++----------- 1 file changed, 27 insertions(+), 35 deletions(-) diff --git a/tests/core/index_management/test_index_management.py b/tests/core/index_management/test_index_management.py index 87d550028..9ab381b0b 100644 --- a/tests/core/index_management/test_index_management.py +++ b/tests/core/index_management/test_index_management.py @@ -100,9 +100,10 @@ def test_bootstrap_vespa_should_successfully_bootstrap_a_new_vespa_application_p self.assertEqual(self.index_management.get_marqo_version(), version.get_version()) + @patch('marqo.vespa.vespa_client.VespaClient.check_for_application_convergence') @patch('marqo.vespa.vespa_client.VespaClient.get_vespa_version') def test_bootstrap_vespa_should_skip_bootstrapping_if_already_bootstrapped_for_older_vespa_version( - self, mock_vespa_version): + self, mock_vespa_version, mock_check_convergence): mock_vespa_version.return_value = '8.382.21' def modified_post(*args, **kwargs): @@ -111,16 +112,31 @@ def modified_post(*args, **kwargs): # verify the first boostrap call deploys the app to vespa with mock.patch.object(httpx.Client, 'post', side_effect=modified_post) as mock_post: self.assertTrue(self.index_management.bootstrap_vespa()) - self.assertEqual(mock_post.call_count, 2) - self.assertTrue('prepareandactivate' in mock_post.call_args_list[1].args[0]) + self.assertEqual(mock_post.call_count, 3) + # First call creates a session to download the app for app version check + self.assertTrue('session?from=' in mock_post.call_args_list[0].args[0]) + # Second call creates a session to download the app to do bootstrapping + self.assertTrue('session?from=' in mock_post.call_args_list[1].args[0]) + # Third call deploys the app by uploading the zip file + self.assertTrue('prepareandactivate' in mock_post.call_args_list[2].args[0]) + + # The first bootstrapping will deploy a new Vespa app, so it will check convergence + mock_check_convergence.assert_called_once() + + mock_check_convergence.reset_mock() # verify the second boostrap call skips the deployment with mock.patch.object(httpx.Client, 'post', side_effect=modified_post) as mock_post: self.assertFalse(self.index_management.bootstrap_vespa()) self.assertEqual(mock_post.call_count, 1) - self.assertFalse('prepareandactivate' in mock_post.call_args_list[0].args[0]) + # First call creates a session to download the app for app version check + self.assertTrue('session?from=' in mock_post.call_args_list[0].args[0]) + + # The second bootstrapping only need to check version, so it will skip convergence check + mock_check_convergence.assert_not_called() - def test_bootstrap_vespa_should_skip_bootstrapping_if_already_bootstrapped(self): + @patch('marqo.vespa.vespa_client.VespaClient.check_for_application_convergence') + def test_bootstrap_vespa_should_skip_bootstrapping_if_already_bootstrapped(self, mock_check_convergence): def modified_put(*args, **kwargs): return httpx.put(*args, **kwargs) @@ -129,11 +145,17 @@ def modified_put(*args, **kwargs): self.assertTrue(self.index_management.bootstrap_vespa()) self.assertTrue('prepare' in mock_post.call_args_list[-2].args[0]) self.assertTrue('active' in mock_post.call_args_list[-1].args[0]) + # The first bootstrapping will deploy a new Vespa app, so it will check convergence + mock_check_convergence.assert_called_once() + + mock_check_convergence.reset_mock() # verify the second boostrap call skips the deployment with mock.patch.object(httpx.Client, 'put', side_effect=modified_put) as mock_post: self.assertFalse(self.index_management.bootstrap_vespa()) self.assertEqual(mock_post.call_count, 0) + # The second bootstrapping only need to check version, so it will skip convergence check + mock_check_convergence.assert_not_called() def test_boostrap_vespa_should_migrate_index_settings_from_existing_vespa_app(self): """ @@ -243,36 +265,6 @@ def test_rollback_should_succeed(self): os.path.join(latest_version, *file) ) - @patch('marqo.vespa.vespa_client.VespaClient.check_for_application_convergence') - def test_bootstrap_and_rollback_should_skip_convergence_check(self, mock_check_convergence): - self.index_management.bootstrap_vespa() - mock_check_convergence.assert_not_called() - - mock_check_convergence.reset_mock() - - try: - self.index_management.rollback_vespa() - except ApplicationRollbackError: - pass - mock_check_convergence.assert_not_called() - - @patch('marqo.vespa.vespa_client.VespaClient.check_for_application_convergence') - @patch('marqo.vespa.vespa_client.VespaClient.get_vespa_version') - def test_bootstrap_and_rollback_should_not_skip_convergence_check_for_older_vespa_version(self, mock_vespa_version, - mock_check_convergence): - mock_vespa_version.return_value = '8.382.21' - - self.index_management.bootstrap_vespa() - mock_check_convergence.assert_called_once() - - mock_check_convergence.reset_mock() - - try: - self.index_management.rollback_vespa() - except ApplicationRollbackError: - pass - mock_check_convergence.assert_called_once() - def test_rollback_should_fail_when_target_version_is_current_version(self): self.index_management.bootstrap_vespa() with self.assertRaisesStrict(ApplicationRollbackError) as e: