1#
2# Copyright 2015 Google Inc.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Various utilities used in tests."""
17
18import contextlib
19import os
20import shutil
21import sys
22import tempfile
23import unittest
24
25import six
26
27
28SkipOnWindows = unittest.skipIf(
29    os.name == 'nt', 'Does not run on windows')
30
31
32@contextlib.contextmanager
33def TempDir(change_to=False):
34    if change_to:
35        original_dir = os.getcwd()
36    path = tempfile.mkdtemp()
37    try:
38        if change_to:
39            os.chdir(path)
40        yield path
41    finally:
42        if change_to:
43            os.chdir(original_dir)
44        shutil.rmtree(path)
45
46
47@contextlib.contextmanager
48def CaptureOutput():
49    new_stdout, new_stderr = six.StringIO(), six.StringIO()
50    old_stdout, old_stderr = sys.stdout, sys.stderr
51    try:
52        sys.stdout, sys.stderr = new_stdout, new_stderr
53        yield new_stdout, new_stderr
54    finally:
55        sys.stdout, sys.stderr = old_stdout, old_stderr
56