diff --git a/main.py b/main.py index c5d63be..1f4f047 100644 --- a/main.py +++ b/main.py @@ -1088,7 +1088,7 @@ def check_api_access(client: httpx.Client, profile_id: str) -> bool: ) elif code == 404: log.critical( - f"{Colors.FAIL}🔍 Profile Not Found: The ID '{profile_id}' does not exist.{Colors.ENDC}" + f"{Colors.FAIL}🔍 Profile Not Found: The ID '{sanitize_for_log(profile_id)}' does not exist.{Colors.ENDC}" ) log.critical( f"{Colors.FAIL} Please verify the Profile ID from your Control D Dashboard URL.{Colors.ENDC}" @@ -1360,11 +1360,15 @@ def delete_folder( ) -> bool: try: _api_delete(client, f"{API_BASE}/{profile_id}/groups/{folder_id}") - log.info("Deleted folder %s (ID %s)", sanitize_for_log(name), folder_id) + log.info( + "Deleted folder %s (ID %s)", + sanitize_for_log(name), + sanitize_for_log(folder_id), + ) return True except httpx.HTTPError as e: log.error( - f"Failed to delete folder {sanitize_for_log(name)} (ID {folder_id}): {sanitize_for_log(e)}" + f"Failed to delete folder {sanitize_for_log(name)} (ID {sanitize_for_log(folder_id)}): {sanitize_for_log(e)}" ) return False @@ -1393,7 +1397,9 @@ def create_folder( if isinstance(body, dict) and "group" in body and "PK" in body["group"]: pk = body["group"]["PK"] log.info( - "Created folder %s (ID %s) [Direct]", sanitize_for_log(name), pk + "Created folder %s (ID %s) [Direct]", + sanitize_for_log(name), + sanitize_for_log(pk), ) return str(pk) @@ -1404,7 +1410,7 @@ def create_folder( log.info( "Created folder %s (ID %s) [Direct]", sanitize_for_log(name), - grp["PK"], + sanitize_for_log(grp["PK"]), ) return str(grp["PK"]) except Exception as e: @@ -1423,7 +1429,7 @@ def create_folder( log.info( "Created folder %s (ID %s) [Polled]", sanitize_for_log(name), - grp["PK"], + sanitize_for_log(grp["PK"]), ) return str(grp["PK"]) except Exception as e: diff --git a/tests/test_plan_details.py b/tests/test_plan_details.py index a7fe865..4f87139 100644 --- a/tests/test_plan_details.py +++ b/tests/test_plan_details.py @@ -28,8 +28,9 @@ def test_print_plan_details_no_colors(capsys): output = captured.out assert "Plan Details for test_profile:" in output - assert " - Folder A: 10 rules" in output - assert " - Folder B: 5 rules" in output + # Match exact output including alignment spaces + assert " - Folder A : 10 rules" in output + assert " - Folder B : 5 rules" in output # Verify alphabetical ordering (A before B) assert output.index("Folder A") < output.index("Folder B") diff --git a/tests/test_push_rules_perf.py b/tests/test_push_rules_perf.py index 120bc29..1019d2e 100644 --- a/tests/test_push_rules_perf.py +++ b/tests/test_push_rules_perf.py @@ -9,10 +9,13 @@ class TestPushRulesPerf(unittest.TestCase): def setUp(self): - # Ensure we are using the current main module instance (in case of reloads) - global main + # Import main freshly or get current from sys.modules + # Because other tests (like test_parallel_deletion.py) might reload main if 'main' in sys.modules: - main = sys.modules['main'] + self.main = sys.modules['main'] + else: + import main + self.main = main self.client = MagicMock() self.profile_id = "test_profile" @@ -22,9 +25,8 @@ def setUp(self): self.status = 1 self.existing_rules = set() - @patch("main.concurrent.futures.as_completed") @patch("main.concurrent.futures.ThreadPoolExecutor") - def test_push_rules_single_batch_optimization(self, mock_executor, mock_as_completed): + def test_push_rules_single_batch_optimization(self, mock_executor): """ Test that push_rules avoids ThreadPoolExecutor for single batch (< 500 rules). """ @@ -41,28 +43,23 @@ def test_push_rules_single_batch_optimization(self, mock_executor, mock_as_compl mock_future.result.return_value = hostnames # Success mock_executor_instance.submit.return_value = mock_future - # Mock as_completed to yield the future immediately - mock_as_completed.return_value = [mock_future] - - # Since we are bypassing TPE, we might need to mock API call? - # The code will call process_batch(1, batch). - # process_batch calls _api_post_form. - # client is mocked, so _api_post_form works (retries mocked). - # But we need to ensure process_batch works correctly in isolation. - - # For this test, we mock _api_post_form? - # No, _api_post_form calls client.post. - - self.main.push_rules( - self.profile_id, - self.folder_name, - self.folder_id, - self.do, - self.status, - hostnames, - self.existing_rules, - self.client - ) + # We also need to mock _api_post_form since it will be called directly + # patch("main._api_post_form") patches what is in sys.modules['main'] + # self.main is likely sys.modules['main'] due to setUp logic + + with patch("main._api_post_form") as mock_post: + self.main.push_rules( + self.profile_id, + self.folder_name, + self.folder_id, + self.do, + self.status, + hostnames, + self.existing_rules, + self.client + ) + + self.assertTrue(mock_post.called, "Expected _api_post_form to be called") # Verify if Executor was called. # AFTER OPTIMIZATION: This should be False. @@ -87,48 +84,51 @@ def test_push_rules_multi_batch(self, mock_executor, mock_as_completed): mock_as_completed.return_value = [mock_future, mock_future] # 2 batches - self.main.push_rules( - self.profile_id, - self.folder_name, - self.folder_id, - self.do, - self.status, - hostnames, - self.existing_rules, - self.client - ) + with patch("main._api_post_form") as mock_post: + self.main.push_rules( + self.profile_id, + self.folder_name, + self.folder_id, + self.do, + self.status, + hostnames, + self.existing_rules, + self.client + ) # This should ALWAYS be True self.assertTrue(mock_executor.called, "ThreadPoolExecutor should be called for multi-batch") - @patch.object(main, "RULE_PATTERN") - def test_push_rules_skips_validation_for_existing(self, mock_rule_pattern): + def test_push_rules_skips_validation_for_existing(self): """ Test that RULE_PATTERN.match is NOT called for rules that are already in existing_rules. """ - # Configure the mock match method - mock_match = mock_rule_pattern.match - mock_match.return_value = True - - hostnames = ["h1", "h2"] - # h1 is already known, h2 is new - existing_rules = {"h1"} - - main.push_rules( - self.profile_id, - self.folder_name, - self.folder_id, - self.do, - self.status, - hostnames, - existing_rules, - self.client - ) - - # h1 is in existing_rules, so we should skip validation for it. - # h2 is NOT in existing_rules, so we should validate it. - # So match should be called EXACTLY once, with "h2". - mock_match.assert_called_once_with("h2") + # Patch RULE_PATTERN on the current main module + with patch.object(self.main, "RULE_PATTERN") as mock_rule_pattern: + # Configure the mock match method + mock_match = mock_rule_pattern.match + mock_match.return_value = True + + hostnames = ["h1", "h2"] + # h1 is already known, h2 is new + existing_rules = {"h1"} + + with patch("main._api_post_form"): + self.main.push_rules( + self.profile_id, + self.folder_name, + self.folder_id, + self.do, + self.status, + hostnames, + existing_rules, + self.client + ) + + # h1 is in existing_rules, so we should skip validation for it. + # h2 is NOT in existing_rules, so we should validate it. + # So match should be called EXACTLY once, with "h2". + mock_match.assert_called_once_with("h2") if __name__ == '__main__': unittest.main()