xref: /petsc/src/binding/petsc4py/test/test_stdout.py (revision 6dd63270497ad23dcf16ae500a87ff2b2a0b7474)
1import unittest
2
3# --------------------------------------------------------------------
4
5
6class TestStdout(unittest.TestCase):
7    def testStdoutRedirect(self):
8        from io import StringIO
9        import sys
10
11        prevstdout = sys.stdout
12        prevstderr = sys.stderr
13        sys.stdout = StringIO()
14        sys.stderr = StringIO()
15
16        import numpy as np
17        from petsc4py import PETSc
18
19        if not (__name__ == '__main__'):
20            PETSc._push_python_vfprintf()
21
22        a = np.array([0.0, 0.0, 0.0], dtype=PETSc.ScalarType)
23        a_vec = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF)
24        a_vec.view()
25        v = PETSc.Viewer.STDERR(PETSc.COMM_SELF)
26        v.printfASCII('Error message')
27
28        newstdout = sys.stdout
29        newstderr = sys.stderr
30        sys.stdout = prevstdout
31        sys.stderr = prevstderr
32
33        output = newstdout.getvalue()
34        error = newstderr.getvalue()
35        if not (__name__ == '__main__'):
36            PETSc._pop_python_vfprintf()
37        stdoutshouldbe = """Vec Object: 1 MPI process
38  type: seq
390.
400.
410.
42"""
43        stderrshouldbe = 'Error message'
44        if PETSc._stdout_is_stderr():
45            stdoutshouldbe = stdoutshouldbe + stderrshouldbe
46            stderrshouldbe = ''
47        self.assertEqual(output, stdoutshouldbe)
48        self.assertEqual(error, stderrshouldbe)
49
50
51# --------------------------------------------------------------------
52
53if __name__ == '__main__':
54    unittest.main()
55
56# --------------------------------------------------------------------
57