diff --git a/src/core/tests/Test_ConfigurePatchingProcessor.py b/src/core/tests/Test_ConfigurePatchingProcessor.py index 76db40e3..11f1d38c 100644 --- a/src/core/tests/Test_ConfigurePatchingProcessor.py +++ b/src/core/tests/Test_ConfigurePatchingProcessor.py @@ -18,7 +18,11 @@ import re import unittest import sys -from io import StringIO +# Conditional import for StringIO +try: + from StringIO import StringIO # Python 2 +except ImportError: + from io import StringIO # Python 3 from core.src.CoreMain import CoreMain from core.src.bootstrap.Constants import Constants @@ -317,6 +321,7 @@ def test_configure_patching_with_patch_mode_and_assessment_mode_by_platform(self def test_configure_patching_raise_exception_auto_os_patch_state(self): # arrange capture std IO captured_output = StringIO() + original_stdout = sys.stdout sys.stdout = captured_output argument_composer = ArgumentComposer() @@ -334,7 +339,7 @@ def test_configure_patching_raise_exception_auto_os_patch_state(self): runtime.configure_patching_processor.start_configure_patching() # restore sdt.out ouptput - sys.stdout = sys.__stdout__ + sys.stdout = original_stdout # assert output = captured_output.getvalue() diff --git a/src/core/tests/Test_MaintenanceWindow.py b/src/core/tests/Test_MaintenanceWindow.py index e9e0fb29..18fe7857 100644 --- a/src/core/tests/Test_MaintenanceWindow.py +++ b/src/core/tests/Test_MaintenanceWindow.py @@ -15,12 +15,17 @@ # Requires Python 2.7+ import datetime -from io import StringIO import sys +# Conditional import for StringIO +try: + from StringIO import StringIO # Python 2 +except ImportError: + from io import StringIO # Python 3 import unittest from core.tests.library.ArgumentComposer import ArgumentComposer from core.tests.library.RuntimeCompositor import RuntimeCompositor + class TestMaintenanceWindow(unittest.TestCase): def setUp(self): pass @@ -67,6 +72,7 @@ def test_RemainingTime_after_duration_complete(self): def test_RemainingTime_log_to_stdout_true(self): # Arrange, Capture stdout captured_output = StringIO() + original_output = sys.stdout sys.stdout = captured_output # Redirect stdout to the StringIO object argument_composer = ArgumentComposer() @@ -78,7 +84,7 @@ def test_RemainingTime_log_to_stdout_true(self): remaining_time = runtime.maintenance_window.get_remaining_time_in_minutes(current_time, log_to_stdout=True) # Restore stdout - sys.stdout = sys.__stdout__ + sys.stdout = original_output # Assert output = captured_output.getvalue()