util/logroll.py

#! /usr/bin/env python3
#
# Copyright 2021 Google Inc.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met: redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer;
# redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution;
# neither the name of the copyright holders nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import copy
import signal
import sys
import unittest
import unittest.mock

parser = argparse.ArgumentParser(
        description='''Circular buffer for text output.

    To capture the rolling last 25 lines of output from command "command":
    command | logroll.py -n 25

    While that's running, to see the most recent 25 lines of output without
    interrupting "command", send SIGUSR1 to the logroll.py process:
    kill -s USR1 ${PID of logroll.py}''')
parser.add_argument('-n', '--lines', default=10, type=int,
        help='Maximum number of lines to buffer at a time.')
parser.add_argument('file', nargs='?', default=sys.stdin,
        type=argparse.FileType('r', encoding='UTF-8'),
        help='File to read from, default is stdin')

args = parser.parse_args()

def dump_lines(lines, idx):
    for line in lines[idx:]:
        print(line, end='')
    for line in lines[:idx]:
        print(line, end='')

def dump_and_exit(lines, idx):
    dump_lines(lines, idx)
    sys.exit(0)

def main(target, incoming):
    idx = 0
    lines = []
    last_idx = target - 1

    signal.signal(signal.SIGUSR1, lambda num, frame: dump_lines(lines, idx))
    signal.signal(signal.SIGINT, lambda num, frame: dump_and_exit(lines, idx))

    for line in incoming:
        lines.append(line)
        if idx == last_idx:
            idx = 0
            break
        else:
            idx += 1

    for lines[idx] in incoming:
        idx = 0 if idx == last_idx else idx + 1

    dump_lines(lines, idx)

if __name__ == '__main__':
    main(target=args.lines, incoming=args.file)



# Unit tests #

class CopyingMock(unittest.mock.MagicMock):
    def __call__(self, *args, **kwargs):
        args = copy.deepcopy(args)
        kwargs = copy.deepcopy(kwargs)
        return super(CopyingMock, self).__call__(*args, **kwargs)

class TestLogroll(unittest.TestCase):
    # Test data.
    lines2 = ['First line',
              'Second line']
    lines3 = ['First line',
              'Second line',
              'Third line']
    lines8 = ['First line',
              'Second line',
              'Third line',
              'Fourth line',
              'Fifth line',
              'Sixth line',
              'Seventh line',
              'Eigth line']

    # Generator which returns lines like a file object would.
    def line_gen(self, lines):
        for line in lines:
            yield line

    # Generator like above, but which simulates a signal midway through.
    def signal_line_gen(self, lines, pos, sig_dict, signal):
        # Return the first few lines.
        for line in lines[:pos]:
            yield line

        # Simulate receiving the signal.
        self.assertIn(signal, sig_dict)
        if signal in sig_dict:
            # Pas in junk for the num and frame arguments.
            sig_dict[signal](None, None)

        # Return the remaining lines.
        for line in lines[pos:]:
            yield line

    # Set up a mock of signal.signal to record handlers in a dict.
    def mock_signal_dict(self, mock):
        signal_dict = {}
        def signal_signal(num, action):
            signal_dict[num] = action
        mock.side_effect = signal_signal
        return signal_dict

    # Actual test methods.
    def test_filling_dump_lines(self):
        with unittest.mock.patch('builtins.print') as mock_print:
            dump_lines(self.lines2, len(self.lines2))
            calls = list([ unittest.mock.call(line, end='') for
                line in self.lines2 ])
            mock_print.assert_has_calls(calls)

    def test_full_dump_lines(self):
        with unittest.mock.patch('builtins.print') as mock_print:
            dump_lines(self.lines2, 0)
            calls = list([ unittest.mock.call(line, end='') for
                line in self.lines2 ])
            mock_print.assert_has_calls(calls)

    def test_offset_dump_lines(self):
        with unittest.mock.patch('builtins.print') as mock_print:
            dump_lines(self.lines3, 1)
            calls = [ unittest.mock.call(self.lines3[1], end=''),
                      unittest.mock.call(self.lines3[2], end=''),
                      unittest.mock.call(self.lines3[0], end='') ]
            mock_print.assert_has_calls(calls)

    def test_dump_and_exit(self):
        with unittest.mock.patch('sys.exit') as mock_sys_exit, \
             unittest.mock.patch(__name__ + '.dump_lines',
                     new_callable=CopyingMock) as mock_dump_lines:
            idx = 1
            dump_and_exit(self.lines3, idx)
            mock_dump_lines.assert_called_with(self.lines3, idx)
            mock_sys_exit.assert_called_with(0)

    def test_filling_main(self):
        with unittest.mock.patch('builtins.print') as mock_print:
            main(5, self.line_gen(self.lines3))
            calls = list([ unittest.mock.call(line, end='') for
                line in self.lines3 ])
            mock_print.assert_has_calls(calls)

    def test_full_main(self):
        with unittest.mock.patch('builtins.print') as mock_print:
            main(5, self.line_gen(self.lines8))
            calls = list([ unittest.mock.call(line, end='') for
                line in self.lines8[-5:] ])
            mock_print.assert_has_calls(calls)

    def test_sigusr1_filling_main(self):
        with unittest.mock.patch('signal.signal') as mock_signal, \
             unittest.mock.patch(__name__ + '.dump_lines',
                     new_callable=CopyingMock) as mock_dump_lines:

            signal_dict = self.mock_signal_dict(mock_signal)

            main(4, self.signal_line_gen(
                self.lines8, 3, signal_dict, signal.SIGUSR1))

            mock_dump_lines.assert_has_calls([
                unittest.mock.call(self.lines8[0:3], 3 % 4),
                unittest.mock.call(self.lines8[-4:], len(self.lines8) % 4)
            ])

    def test_sigint_filling_main(self):
        with unittest.mock.patch('signal.signal') as mock_signal, \
             unittest.mock.patch(__name__ + '.dump_lines',
                     new_callable=CopyingMock) as mock_dump_lines:

            signal_dict = self.mock_signal_dict(mock_signal)

            with self.assertRaises(SystemExit):
                main(4, self.signal_line_gen(
                    self.lines8, 3, signal_dict, signal.SIGINT))

            mock_dump_lines.assert_has_calls([
                unittest.mock.call(self.lines8[0:3], 3 % 4),
            ])

    def test_sigusr1_full_main(self):
        with unittest.mock.patch('signal.signal') as mock_signal, \
             unittest.mock.patch(__name__ + '.dump_lines',
                     new_callable=CopyingMock) as mock_dump_lines:

            signal_dict = self.mock_signal_dict(mock_signal)

            main(4, self.signal_line_gen(
                self.lines8, 5, signal_dict, signal.SIGUSR1))

            mock_dump_lines.assert_has_calls([
                unittest.mock.call(self.lines8[4:5] + self.lines8[1:4], 5 % 4),
                unittest.mock.call(self.lines8[-4:], len(self.lines8) % 4)
            ])

    def test_sigint_full_main(self):
        with unittest.mock.patch('signal.signal') as mock_signal, \
             unittest.mock.patch(__name__ + '.dump_lines',
                     new_callable=CopyingMock) as mock_dump_lines:

            signal_dict = self.mock_signal_dict(mock_signal)

            with self.assertRaises(SystemExit):
                main(4, self.signal_line_gen(
                    self.lines8, 5, signal_dict, signal.SIGINT))

            mock_dump_lines.assert_has_calls([
                unittest.mock.call(self.lines8[4:5] + self.lines8[1:4], 5 % 4),
            ])