from docutils import nodes
from sphinx.util.compat import Directive
import unittest


class RecordingTestResult(unittest.TestResult):
    def __init__(self):
        super().__init__()
        self.successfull_tests = []
        self.failed_tests = []
        self.skipped_tests = []

    def addSuccess(self, tst):
        super().addSuccess(tst)
        self.successfull_tests.append(tst)

    def addFailure(self, tst, err):
        super().addFailure(tst, err)
        self.failed_tests.append(tst)
    def addError(self, tst, err):
        super().addError(tst, err)
        self.failed_tests.append(tst)
    def addSkip(self, tst, reason):
        super().addSkip(tst, reason)
        self.skipped_tests.append(tst)

def makeTable(rows):
    # Buildup table:
    table = nodes.table()
    group = nodes.tgroup()
    group += nodes.colspec('', colwidth=100)
    group += nodes.colspec('', colwidth=50)
    table += group
    tbody = nodes.tbody('')
    group += tbody
    for c1, c2 in rows:
        row = nodes.row('')
        row += nodes.entry('', nodes.paragraph('', c1))
        row += nodes.entry('', nodes.paragraph('', c2))
        tbody += row
    return table

class TestResults(Directive):
    required_arguments = 1
    optional_arguments = 0
    option_spec = {}
    has_content = False
    node_class = None

    def run(self):
        # Run unit test:
        tl = unittest.TestLoader()
        result = RecordingTestResult()
        tests = tl.discover('../test')
        tests.run(result)
        para = nodes.paragraph()
        para += nodes.Text('Unit test results:')
        res = []
        for s in result.successfull_tests:
            res.append((str(s), 'OK'))
        for s in result.skipped_tests:
            res.append((str(s), 'SKIP'))
        for s in result.failed_tests:
            res.append((str(s), 'FAIL'))
        table = makeTable(res)

        return [para, table]

def setup(app):
    app.add_directive('testresults', TestResults)