# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

"""
Test for distributed trial worker side.
"""

import os
from io import BytesIO, StringIO

from zope.interface.verify import verifyObject

from twisted.internet.defer import fail, succeed
from twisted.internet.error import ProcessDone
from twisted.internet.interfaces import IAddress, ITransport
from twisted.protocols.amp import AMP
from twisted.python.failure import Failure
from twisted.python.filepath import FilePath
from twisted.python.reflect import fullyQualifiedName
from twisted.scripts import trial
from twisted.test.proto_helpers import StringTransport
from twisted.trial._dist import managercommands, workercommands
from twisted.trial._dist.worker import (
    LocalWorker,
    LocalWorkerAMP,
    LocalWorkerTransport,
    NotRunning,
    WorkerException,
    WorkerProtocol,
)
from twisted.trial.reporter import TestResult
from twisted.trial.unittest import TestCase


class FakeAMP(AMP):
    """
    A fake amp protocol.
    """


class WorkerProtocolTests(TestCase):
    """
    Tests for L{WorkerProtocol}.
    """

    def setUp(self):
        """
        Set up a transport, a result stream and a protocol instance.
        """
        self.serverTransport = StringTransport()
        self.clientTransport = StringTransport()
        self.server = WorkerProtocol()
        self.server.makeConnection(self.serverTransport)
        self.client = FakeAMP()
        self.client.makeConnection(self.clientTransport)

    def test_run(self):
        """
        Calling the L{workercommands.Run} command on the client returns a
        response with C{success} sets to C{True}.
        """
        d = self.client.callRemote(workercommands.Run, testCase="doesntexist")

        def check(result):
            self.assertTrue(result["success"])

        d.addCallback(check)
        self.server.dataReceived(self.clientTransport.value())
        self.clientTransport.clear()
        self.client.dataReceived(self.serverTransport.value())
        self.serverTransport.clear()
        return d

    def test_start(self):
        """
        The C{start} command changes the current path.
        """
        curdir = os.path.realpath(os.path.curdir)
        self.addCleanup(os.chdir, curdir)
        self.server.start("..")
        self.assertNotEqual(os.path.realpath(os.path.curdir), curdir)


class LocalWorkerAMPTests(TestCase):
    """
    Test case for distributed trial's manager-side local worker AMP protocol
    """

    def setUp(self):
        self.managerTransport = StringTransport()
        self.managerAMP = LocalWorkerAMP()
        self.managerAMP.makeConnection(self.managerTransport)
        self.result = TestResult()
        self.workerTransport = StringTransport()
        self.worker = AMP()
        self.worker.makeConnection(self.workerTransport)

        config = trial.Options()
        self.testName = "twisted.doesnexist"
        config["tests"].append(self.testName)
        self.testCase = trial._getSuite(config)._tests.pop()

        self.managerAMP.run(self.testCase, self.result)
        self.managerTransport.clear()

    def pumpTransports(self):
        """
        Sends data from C{self.workerTransport} to C{self.managerAMP}, and then
        data from C{self.managerTransport} back to C{self.worker}.
        """
        self.managerAMP.dataReceived(self.workerTransport.value())
        self.workerTransport.clear()
        self.worker.dataReceived(self.managerTransport.value())

    def test_runSuccess(self):
        """
        Run a test, and succeed.
        """
        results = []

        d = self.worker.callRemote(managercommands.AddSuccess, testName=self.testName)
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        self.assertTrue(results)

    def test_runExpectedFailure(self):
        """
        Run a test, and fail expectedly.
        """
        results = []

        d = self.worker.callRemote(
            managercommands.AddExpectedFailure,
            testName=self.testName,
            error="error",
            todo="todoReason",
        )
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        self.assertEqual(self.testCase, self.result.expectedFailures[0][0])
        self.assertTrue(results)

    def test_runError(self):
        """
        Run a test, and encounter an error.
        """
        results = []
        errorClass = fullyQualifiedName(ValueError)
        d = self.worker.callRemote(
            managercommands.AddError,
            testName=self.testName,
            error="error",
            errorClass=errorClass,
            frames=[],
        )
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        case, failure = self.result.errors[0]
        self.assertEqual(self.testCase, case)
        self.assertEqual(failure.type, ValueError)
        self.assertEqual(failure.value, WorkerException("error"))
        self.assertTrue(results)

    def test_runErrorWithFrames(self):
        """
        L{LocalWorkerAMP._buildFailure} recreates the C{Failure.frames} from
        the C{frames} argument passed to C{AddError}.
        """
        results = []
        errorClass = fullyQualifiedName(ValueError)
        d = self.worker.callRemote(
            managercommands.AddError,
            testName=self.testName,
            error="error",
            errorClass=errorClass,
            frames=["file.py", "invalid code", "3"],
        )
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        case, failure = self.result.errors[0]
        self.assertEqual(self.testCase, case)
        self.assertEqual(failure.type, ValueError)
        self.assertEqual(failure.value, WorkerException("error"))
        self.assertEqual([("file.py", "invalid code", 3, [], [])], failure.frames)
        self.assertTrue(results)

    def test_runFailure(self):
        """
        Run a test, and fail.
        """
        results = []
        failClass = fullyQualifiedName(RuntimeError)
        d = self.worker.callRemote(
            managercommands.AddFailure,
            testName=self.testName,
            fail="fail",
            failClass=failClass,
            frames=[],
        )
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        case, failure = self.result.failures[0]
        self.assertEqual(self.testCase, case)
        self.assertEqual(failure.type, RuntimeError)
        self.assertEqual(failure.value, WorkerException("fail"))
        self.assertTrue(results)

    def test_runSkip(self):
        """
        Run a test, but skip it.
        """
        results = []

        d = self.worker.callRemote(
            managercommands.AddSkip, testName=self.testName, reason="reason"
        )
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        self.assertEqual(self.testCase, self.result.skips[0][0])
        self.assertTrue(results)

    def test_runUnexpectedSuccesses(self):
        """
        Run a test, and succeed unexpectedly.
        """
        results = []

        d = self.worker.callRemote(
            managercommands.AddUnexpectedSuccess, testName=self.testName, todo="todo"
        )
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        self.assertEqual(self.testCase, self.result.unexpectedSuccesses[0][0])
        self.assertTrue(results)

    def test_testWrite(self):
        """
        L{LocalWorkerAMP.testWrite} writes the data received to its test
        stream.
        """
        results = []
        stream = StringIO()
        self.managerAMP.setTestStream(stream)

        command = managercommands.TestWrite
        d = self.worker.callRemote(command, out="Some output")
        d.addCallback(lambda result: results.append(result["success"]))
        self.pumpTransports()

        self.assertEqual("Some output\n", stream.getvalue())
        self.assertTrue(results)

    def test_stopAfterRun(self):
        """
        L{LocalWorkerAMP.run} calls C{stopTest} on its test result once the
        C{Run} commands has succeeded.
        """
        result = object()
        stopped = []

        def fakeCallRemote(command, testCase):
            return succeed(result)

        self.managerAMP.callRemote = fakeCallRemote

        class StopTestResult(TestResult):
            def stopTest(self, test):
                stopped.append(test)

        d = self.managerAMP.run(self.testCase, StopTestResult())
        self.assertEqual([self.testCase], stopped)
        return d.addCallback(self.assertIdentical, result)


class SpyDataLocalWorkerAMP(LocalWorkerAMP):
    """
    A fake implementation of L{LocalWorkerAMP} that records the received
    data and doesn't automatically dispatch any command..
    """

    id = 0
    dataString = b""

    def dataReceived(self, data):
        self.dataString += data


class FakeTransport:
    """
    A fake process transport implementation for testing.
    """

    dataString = b""
    calls = 0

    def writeToChild(self, fd, data):
        self.dataString += data

    def loseConnection(self):
        self.calls += 1


class LocalWorkerTests(TestCase):
    """
    Tests for L{LocalWorker} and L{LocalWorkerTransport}.
    """

    def tidyLocalWorker(self, *args, **kwargs):
        """
        Create a L{LocalWorker}, connect it to a transport, and ensure
        its log files are closed.

        @param args: See L{LocalWorker}

        @param kwargs: See L{LocalWorker}

        @return: a L{LocalWorker} instance
        """
        worker = LocalWorker(*args, **kwargs)
        worker.makeConnection(FakeTransport())
        self.addCleanup(worker._outLog.close)
        self.addCleanup(worker._errLog.close)
        return worker

    def test_exitBeforeConnected(self):
        """
        L{LocalWorker.exit} fails with L{NotRunning} if it is called before the
        protocol is connected to a transport.
        """
        worker = LocalWorker(
            SpyDataLocalWorkerAMP(), FilePath(self.mktemp()), StringIO()
        )
        self.failureResultOf(worker.exit(), NotRunning)

    def test_exitAfterDisconnected(self):
        """
        L{LocalWorker.exit} fails with L{NotRunning} if it is called after the the
        protocol is disconnected from its transport.
        """
        worker = self.tidyLocalWorker(
            SpyDataLocalWorkerAMP(), FilePath(self.mktemp()), StringIO()
        )
        worker.processEnded(Failure(ProcessDone(0)))
        # Since we're not calling exit until after the process has ended, it
        # won't consume the ProcessDone failure on the internal `endDeferred`.
        # Swallow it here.
        self.failureResultOf(worker.endDeferred, ProcessDone)

        # Now assert that exit behaves.
        self.failureResultOf(worker.exit(), NotRunning)

    def test_childDataReceived(self):
        """
        L{LocalWorker.childDataReceived} forwards the received data to linked
        L{AMP} protocol if the right file descriptor, otherwise forwards to
        C{ProcessProtocol.childDataReceived}.
        """
        localWorker = self.tidyLocalWorker(
            SpyDataLocalWorkerAMP(), FilePath(self.mktemp()), "test.log"
        )
        localWorker._outLog = BytesIO()
        localWorker.childDataReceived(4, b"foo")
        localWorker.childDataReceived(1, b"bar")
        self.assertEqual(b"foo", localWorker._ampProtocol.dataString)
        self.assertEqual(b"bar", localWorker._outLog.getvalue())

    def test_newlineStyle(self):
        """
        L{LocalWorker} writes the log data with local newlines.
        """
        amp = SpyDataLocalWorkerAMP()
        tempDir = FilePath(self.mktemp())
        tempDir.makedirs()
        logPath = tempDir.child("test.log")

        with open(logPath.path, "wt", encoding="utf-8") as logFile:
            worker = LocalWorker(amp, tempDir, logFile)
            worker.makeConnection(FakeTransport())
            self.addCleanup(worker._outLog.close)
            self.addCleanup(worker._errLog.close)

            expected = "Here comes the \N{sun}!"
            amp.testWrite(expected)

        self.assertEqual(
            # os.linesep is the local newline.
            (expected + os.linesep),
            # getContent reads in binary mode so we'll see the bytes that
            # actually ended up in the file.
            logPath.getContent().decode("utf-8"),
        )

    def test_outReceived(self):
        """
        L{LocalWorker.outReceived} logs the output into its C{_outLog} log
        file.
        """
        localWorker = self.tidyLocalWorker(
            SpyDataLocalWorkerAMP(), FilePath(self.mktemp()), "test.log"
        )
        localWorker._outLog = BytesIO()
        data = b"The quick brown fox jumps over the lazy dog"
        localWorker.outReceived(data)
        self.assertEqual(data, localWorker._outLog.getvalue())

    def test_errReceived(self):
        """
        L{LocalWorker.errReceived} logs the errors into its C{_errLog} log
        file.
        """
        localWorker = self.tidyLocalWorker(
            SpyDataLocalWorkerAMP(), FilePath(self.mktemp()), "test.log"
        )
        localWorker._errLog = BytesIO()
        data = b"The quick brown fox jumps over the lazy dog"
        localWorker.errReceived(data)
        self.assertEqual(data, localWorker._errLog.getvalue())

    def test_write(self):
        """
        L{LocalWorkerTransport.write} forwards the written data to the given
        transport.
        """
        transport = FakeTransport()
        localTransport = LocalWorkerTransport(transport)
        data = b"The quick brown fox jumps over the lazy dog"
        localTransport.write(data)
        self.assertEqual(data, transport.dataString)

    def test_writeSequence(self):
        """
        L{LocalWorkerTransport.writeSequence} forwards the written data to the
        given transport.
        """
        transport = FakeTransport()
        localTransport = LocalWorkerTransport(transport)
        data = (b"The quick ", b"brown fox jumps ", b"over the lazy dog")
        localTransport.writeSequence(data)
        self.assertEqual(b"".join(data), transport.dataString)

    def test_loseConnection(self):
        """
        L{LocalWorkerTransport.loseConnection} forwards the call to the given
        transport.
        """
        transport = FakeTransport()
        localTransport = LocalWorkerTransport(transport)
        localTransport.loseConnection()

        self.assertEqual(transport.calls, 1)

    def test_connectionLost(self):
        """
        L{LocalWorker.connectionLost} closes the per-worker log streams.
        """

        localWorker = self.tidyLocalWorker(
            SpyDataLocalWorkerAMP(), FilePath(self.mktemp()), "test.log"
        )
        localWorker.connectionLost(None)
        self.assertTrue(localWorker._outLog.closed)
        self.assertTrue(localWorker._errLog.closed)

    def test_processEnded(self):
        """
        L{LocalWorker.processEnded} calls C{connectionLost} on itself and on
        the L{AMP} protocol.
        """
        transport = FakeTransport()
        protocol = SpyDataLocalWorkerAMP()
        localWorker = LocalWorker(protocol, FilePath(self.mktemp()), "test.log")
        localWorker.makeConnection(transport)
        localWorker.processEnded(Failure(ProcessDone(0)))
        self.assertTrue(localWorker._outLog.closed)
        self.assertTrue(localWorker._errLog.closed)
        self.assertIdentical(None, protocol.transport)
        return self.assertFailure(localWorker.endDeferred, ProcessDone)

    def test_addresses(self):
        """
        L{LocalWorkerTransport.getPeer} and L{LocalWorkerTransport.getHost}
        return L{IAddress} objects.
        """
        localTransport = LocalWorkerTransport(None)
        self.assertTrue(verifyObject(IAddress, localTransport.getPeer()))
        self.assertTrue(verifyObject(IAddress, localTransport.getHost()))

    def test_transport(self):
        """
        L{LocalWorkerTransport} implements L{ITransport} to be able to be used
        by L{AMP}.
        """
        localTransport = LocalWorkerTransport(None)
        self.assertTrue(verifyObject(ITransport, localTransport))

    def test_startError(self):
        """
        L{LocalWorker} swallows the exceptions returned by the L{AMP} protocol
        start method, as it generates unnecessary errors.
        """

        def failCallRemote(command, directory):
            return fail(RuntimeError("oops"))

        protocol = SpyDataLocalWorkerAMP()
        protocol.callRemote = failCallRemote
        self.tidyLocalWorker(protocol, FilePath(self.mktemp()), "test.log")
        self.assertEqual([], self.flushLoggedErrors(RuntimeError))
