infix/test/infamy/tap.py

"""TAP (Test Anything Protocol) output helpers for Infamy tests.

TAP is a simple line-based protocol for communicating test results between a
test producer and a consumer (harness).  The canonical reference is:

  https://testanything.org/tap-specification.html

The key elements are:

  Plan line:    ``1..N``           — total number of tests, written once
  Result lines: ``ok N - desc``    — test N passed
                ``not ok N - desc``— test N failed
  Directives:   ``ok N # SKIP``    ``not ok N # TODO``
  Diagnostics:  Any line beginning with ``#`` is a comment/diagnostic.

The plan line may appear either at the beginning (before any results) or at
the end (after all results).  Infamy writes it at the END so that the total
step count is known; a harness that requires an upfront plan will report
"no plan" on failure — see CommentWriter below for a known limitation.

``CommentWriter`` redirects Python's ``sys.stdout`` so that ordinary
``print()`` calls are prefixed with ``# `` and appear as diagnostics, while
TAP result lines (written via ``self.out``, the *original* stdout) are emitted
verbatim.  The plan line is also written via ``self.out``.

Note: because ``sys.stdout`` is replaced *after* module import the default
argument ``output=sys.stdout`` captures the original stream at import time,
so nested or re-entrant ``Test()`` instances will share the same ``self.out``.
"""
import contextlib
import datetime
import subprocess
import sys
import traceback

import infamy.netns

class Test:
    def __init__(self, output=sys.stdout):
        self.out = output
        self.commenter = CommentWriter(self.out)
        if self.out == sys.stdout:
            sys.stdout = self.commenter

        self.test_cleanup=[]
        self.steps = 0

    def push_test_cleanup(self, fn):
        self.test_cleanup.append(fn)

    def cleanup(self):
        infamy.netns.IsolatedMacVlans.Cleanup()
        for test_cleanup in reversed(self.test_cleanup):
            test_cleanup()

    def __enter__(self):
        now = datetime.datetime.now().strftime("%F %T")
        self.out.write(f"# Starting ({now})\n")
        self.out.flush()
        return self

    def __exit__(self, t, e, tb):
        now = datetime.datetime.now().strftime("%F %T")
        self.out.write(f"# Exiting ({now})\n")
        self.out.flush()

        self.cleanup()

        if not e:
            self._not_ok("Missing explicit test result\n")
        else:
            if t in (TestPass, TestSkip):
                self.out.write(f"1..{self.steps}\n")
                self.out.flush()
                raise SystemExit(0)
            if t is AssertionError:
                traceback.print_tb(tb, file=self.commenter)
                self.out.write(f"{str(e)}\n")
                self.out.flush()
            else:
                traceback.print_exception(e, file=self.commenter)

            if type(e) is subprocess.CalledProcessError:
                print("Failing subprocess stdout:\n", e.stdout)
            elif len(e.args) and type(e.args[0]) is subprocess.CompletedProcess:
                print("Failing subprocess stdout:\n", e.args[0].stdout)

        self.out.write(f"1..{self.steps}\n")
        self.out.flush()
        raise SystemExit(1)

    @contextlib.contextmanager
    def step(self, msg):
        try:
            yield
            self._ok(msg)

        except Exception as e:
            if type(e) == TestPass:
                self._ok(msg)
            elif type(e) == TestSkip:
                self._ok(directive="skip")
            elif type(e) == TestFail:
                self._not_ok(msg)
            else:
                self._not_ok(msg)

            raise e

    def _report(self, tag, msg):
        self.steps += 1
        self.out.write(f"{tag} {self.steps}{msg}\n")
        self.out.flush()

    def _ok(self, msg="", directive=None):
        if msg:
            msg = " - " + msg

        if directive:
            msg = msg + " # " + directive

        self._report("ok", msg)

    def _not_ok(self, msg):
        self._report("not ok", " - " + msg)

    def succeed(self):
        raise TestPass()

    def skip(self):
        raise TestSkip()

    def fail(self, message=None):
        if message:
            self.out.write(f"# {message}\n")
            self.out.flush()
        raise TestFail()


class TestResult(Exception):
    pass


class TestPass(TestResult):
    pass


class TestFail(TestResult):
    pass


class TestSkip(TestResult):
    pass


class CommentWriter:
    def __init__(self, f):
        self.f = f
        self.at_nl = True

    def write(self, data):
        if self.at_nl:
            data = "# " + data
            self.at_nl = False
            if not len(data):
                return

        if data.endswith("\n"):
            self.at_nl = True
            data = data[:-1]

        data = data.replace("\n", "\n# ")
        if self.at_nl:
            data = data + "\n"

        self.f.write(data)
        self.flush()

    def flush(self):
        return self.f.flush()