22 lines
651 B
Python
22 lines
651 B
Python
from psh.commands import BaseCommand
|
|
|
|
from io import StringIO
|
|
|
|
class TestFormatter(BaseCommand):
|
|
"""Formatter useful for tests. Instead of printing to stdout, it
|
|
stores any output inside a stringio buffer. This can be retrieved
|
|
with the get_data method."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(TestFormatter, self).__init__(*args, **kwargs)
|
|
self.buffer = StringIO()
|
|
|
|
def call(self):
|
|
input_generator = self.get_input_generator()
|
|
for line in input_generator:
|
|
self.buffer.write(line.decode('utf-8'))
|
|
return None
|
|
|
|
def get_data(self):
|
|
return self.buffer.getvalue()
|