From 5b0591392f3ce758b97368333dbaf7eb26151c30 Mon Sep 17 00:00:00 2001
From: Alinson S Xavier <axavier@anl.gov>
Date: Sun, 12 Apr 2020 08:15:23 -0500
Subject: [PATCH] Fix RedirectOutput; add tests

---
 src/python/miplearn/solvers/__init__.py       | 22 +++++++++----------
 .../solvers/tests/test_internal_solver.py     | 12 ++++++++++
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/python/miplearn/solvers/__init__.py b/src/python/miplearn/solvers/__init__.py
index 93ee6d0..ae0d141 100644
--- a/src/python/miplearn/solvers/__init__.py
+++ b/src/python/miplearn/solvers/__init__.py
@@ -2,20 +2,15 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
 import sys
 
+logger = logging.getLogger(__name__)
 
-class RedirectOutput(object):
+
+class RedirectOutput:
     def __init__(self, streams):
         self.streams = streams
-        self._original_stdout = sys.stdout
-        self._original_stderr = sys.stderr
-        sys.stdout = self
-        sys.stderr = self
-
-    def __del__(self):
-        sys.stdout = self._original_stdout
-        sys.stderr = self._original_stderr
 
     def write(self, data):
         for stream in self.streams:
@@ -26,7 +21,12 @@ class RedirectOutput(object):
             stream.flush()
 
     def __enter__(self):
-        pass
+        self._original_stdout = sys.stdout
+        self._original_stderr = sys.stderr
+        sys.stdout = self
+        sys.stderr = self
+        return self
 
     def __exit__(self, _type, _value, _traceback):
-        pass
+        sys.stdout = self._original_stdout
+        sys.stderr = self._original_stderr
diff --git a/src/python/miplearn/solvers/tests/test_internal_solver.py b/src/python/miplearn/solvers/tests/test_internal_solver.py
index 1ee0d12..983adc4 100644
--- a/src/python/miplearn/solvers/tests/test_internal_solver.py
+++ b/src/python/miplearn/solvers/tests/test_internal_solver.py
@@ -1,8 +1,10 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
+from io import StringIO
 
 import pyomo.environ as pe
+from miplearn.solvers import RedirectOutput
 from miplearn.solvers.cplex import CPLEXSolver
 from miplearn.solvers.gurobi import GurobiSolver
 
@@ -10,6 +12,16 @@ from . import _get_instance
 from ...problems.knapsack import ChallengeA
 
 
+def test_redirect_output():
+    import sys
+    original_stdout = sys.stdout
+    io = StringIO()
+    with RedirectOutput([io]):
+        print("Hello world")
+    assert sys.stdout == original_stdout
+    assert io.getvalue() == "Hello world\n"
+
+
 def test_internal_solver_warm_starts():
     for solver in [GurobiSolver(), CPLEXSolver()]:
         instance = _get_instance()