from __future__ import absolute_import, print_function

import unittest

from mercurial import (

    ui as uimod,

    util,

    wireprotoframing as framing,

)

from mercurial.utils import cborutil

ffs = framing.makeframefromhumanstring

OK = b''.join(cborutil.streamencode({b'status': b'ok'}))

def makereactor(deferoutput=False):

    ui = uimod.ui()

    return framing.serverreactor(ui, deferoutput=deferoutput)

def sendframes(reactor, gen):

    """Send a generator of frame bytearray to a reactor.

    Emits a generator of results from ``onframerecv()`` calls.

    """

    for frame in gen:

        header = framing.parseheader(frame)

        payload = frame[framing.FRAME_HEADER_SIZE :]

        assert len(payload) == header.length

        yield reactor.onframerecv(

            framing.frame(

                header.requestid,

                header.streamid,

                header.streamflags,

                header.typeid,

                header.flags,

                payload,

            )

        )

def sendcommandframes(reactor, stream, rid, cmd, args, datafh=None):

    """Generate frames to run a command and send them to a reactor."""

    return sendframes(

        reactor, framing.createcommandframes(stream, rid, cmd, args, datafh)

    )

class ServerReactorTests(unittest.TestCase):

    def _sendsingleframe(self, reactor, f):

        results = list(sendframes(reactor, [f]))

        self.assertEqual(len(results), 1)

        return results[0]

    def assertaction(self, res, expected):

        self.assertIsInstance(res, tuple)

        self.assertEqual(len(res), 2)

        self.assertIsInstance(res[1], dict)

        self.assertEqual(res[0], expected)

    def assertframesequal(self, frames, framestrings):

        expected = [ffs(s) for s in framestrings]

        self.assertEqual(list(frames), expected)

    def test1framecommand(self):

        """Receiving a command in a single frame yields request to run it."""

        reactor = makereactor()

        stream = framing.stream(1)

        results = list(sendcommandframes(reactor, stream, 1, b'mycommand', {}))

        self.assertEqual(len(results), 1)

        self.assertaction(results[0], b'runcommand')

        self.assertEqual(

            results[0][1],

            {

                b'requestid': 1,

                b'command': b'mycommand',

                b'args': {},

                b'redirect': None,

                b'data': None,

            },

        )

        result = reactor.oninputeof()

        self.assertaction(result, b'noop')

    def test1argument(self):

        reactor = makereactor()

        stream = framing.stream(1)

        results = list(

            sendcommandframes(

                reactor, stream, 41, b'mycommand', {b'foo': b'bar'}

            )

        )

        self.assertEqual(len(results), 1)

        self.assertaction(results[0], b'runcommand')

        self.assertEqual(

            results[0][1],

            {

                b'requestid': 41,

                b'command': b'mycommand',

                b'args': {b'foo': b'bar'},

                b'redirect': None,

                b'data': None,

            },

        )

    def testmultiarguments(self):

        reactor = makereactor()

        stream = framing.stream(1)

        results = list(

            sendcommandframes(

                reactor,

                stream,

                1,

                b'mycommand',

                {b'foo': b'bar', b'biz': b'baz'},

            )

        )

        self.assertEqual(len(results), 1)

        self.assertaction(results[0], b'runcommand')

        self.assertEqual(

            results[0][1],

            {

                b'requestid': 1,

                b'command': b'mycommand',

                b'args': {b'foo': b'bar', b'biz': b'baz'},

                b'redirect': None,

                b'data': None,

            },

        )

    def testsimplecommanddata(self):

        reactor = makereactor()

        stream = framing.stream(1)

        results = list(

            sendcommandframes(

                reactor, stream, 1, b'mycommand', {}, util.bytesio(b'data!')

            )

        )

        self.assertEqual(len(results), 2)

        self.assertaction(results[0], b'wantframe')

        self.assertaction(results[1], b'runcommand')

        self.assertEqual(

            results[1][1],

            {

                b'requestid': 1,

                b'command': b'mycommand',

                b'args': {},

                b'redirect': None,

                b'data': b'data!',

            },

        )

    def testmultipledataframes(self):

        frames = [

            ffs(

                b'1 1 stream-begin command-request new|have-data '

                b"cbor:{b'name': b'mycommand'}"

            ),

            ffs(b'1 1 0 command-data continuation data1'),

            ffs(b'1 1 0 command-data continuation data2'),

            ffs(b'1 1 0 command-data eos data3'),

        ]

        reactor = makereactor()

        results = list(sendframes(reactor, frames))

        self.assertEqual(len(results), 4)

        for i in range(3):

            self.assertaction(results[i], b'wantframe')

        self.assertaction(results[3], b'runcommand')

        self.assertEqual(

            results[3][1],

            {

                b'requestid': 1,

                b'command': b'mycommand',

                b'args': {},

                b'redirect': None,

                b'data': b'data1data2data3',

            },

        )

    def testargumentanddata(self):

        frames = [

            ffs(

                b'1 1 stream-begin command-request new|have-data '

                b"cbor:{b'name': b'command', b'args': {b'key': b'val',"

                b"b'foo': b'bar'}}"

            ),

            ffs(b'1 1 0 command-data continuation value1'),

            ffs(b'1 1 0 command-data eos value2'),

        ]

        reactor = makereactor()

        results = list(sendframes(reactor, frames))

        self.assertaction(results[-1], b'runcommand')

        self.assertEqual(

            results[-1][1],

            {

                b'requestid': 1,

                b'command': b'command',

                b'args': {

                    b'key': b'val',

                    b'foo': b'bar',

                },

                b'redirect': None,

                b'data': b'value1value2',

            },

        )

    def testnewandcontinuation(self):

        result = self._sendsingleframe(

            makereactor(),

            ffs(b'1 1 stream-begin command-request new|continuation '),

        )

        self.assertaction(result, b'error')

        self.assertEqual(

            result[1],

            {

                b'message': b'received command request frame with both new and '

                b'continuation flags set',

            },

        )

    def testneithernewnorcontinuation(self):

        result = self._sendsingleframe(

            makereactor(), ffs(b'1 1 stream-begin command-request 0 ')

        )

        self.assertaction(result, b'error')

        self.assertEqual(

            result[1],

            {

                b'message': b'received command request frame with neither new nor '

                b'continuation flags set',

            },

        )

    def testunexpectedcommanddata(self):

        """Command data frame when not running a command is an error."""

        result = self._sendsingleframe(

            makereactor(), ffs(b'1 1 stream-begin command-data 0 ignored')

        )

        self.assertaction(result, b'error')

        self.assertEqual(

            result[1],

            {

                b'message': b'expected sender protocol settings or command request '

                b'frame; got 2',

            },

        )

    def testunexpectedcommanddatareceiving(self):

        """Same as above except the command is receiving."""

        results = list(

            sendframes(

                makereactor(),

                [

                    ffs(

                        b'1 1 stream-begin command-request new|more '

                        b"cbor:{b'name': b'ignored'}"

                    ),

                    ffs(b'1 1 0 command-data eos ignored'),

                ],

            )

        )

        self.assertaction(results[0], b'wantframe')

        self.assertaction(results[1], b'error')

        self.assertEqual(

            results[1][1],

            {

                b'message': b'received command data frame for request that is not '

                b'expecting data: 1',

            },

        )

    def testconflictingrequestidallowed(self):

        """Multiple fully serviced commands with same request ID is allowed."""

        reactor = makereactor()

        results = []

        outstream = reactor.makeoutputstream()

        results.append(

            self._sendsingleframe(

                reactor,

                ffs(

                    b'1 1 stream-begin command-request new '

                    b"cbor:{b'name': b'command'}"

                ),

            )

        )

        result = reactor.oncommandresponsereadyobjects(

            outstream, 1, [b'response1']

        )

        self.assertaction(result, b'sendframes')

        list(result[1][b'framegen'])

        results.append(

            self._sendsingleframe(

                reactor,

                ffs(

                    b'1 1 stream-begin command-request new '

                    b"cbor:{b'name': b'command'}"

                ),

            )

        )

        result = reactor.oncommandresponsereadyobjects(

            outstream, 1, [b'response2']

        )

        self.assertaction(result, b'sendframes')

        list(result[1][b'framegen'])

        results.append(

            self._sendsingleframe(

                reactor,

                ffs(

                    b'1 1 stream-begin command-request new '

                    b"cbor:{b'name': b'command'}"

                ),

            )

        )

        result = reactor.oncommandresponsereadyobjects(

            outstream, 1, [b'response3']

        )

        self.assertaction(result, b'sendframes')

        list(result[1][b'framegen'])

        for i in range(3):

            self.assertaction(results[i], b'runcommand')

            self.assertEqual(

                results[i][1],

                {

                    b'requestid': 1,

                    b'command': b'command',

                    b'args': {},

                    b'redirect': None,

                    b'data': None,

                },

            )

    def testconflictingrequestid(self):

        """Request ID for new command matching in-flight command is illegal."""

        results = list(

            sendframes(

                makereactor(),

                [

                    ffs(

                        b'1 1 stream-begin command-request new|more '

                        b"cbor:{b'name': b'command'}"

                    ),

                    ffs(

                        b'1 1 0 command-request new '

                        b"cbor:{b'name': b'command1'}"

                    ),

                ],

            )

        )

        self.assertaction(results[0], b'wantframe')

        self.assertaction(results[1], b'error')

        self.assertEqual(

            results[1][1],

            {

                b'message': b'request with ID 1 already received',

            },

        )

    def testinterleavedcommands(self):

        cbor1 = b''.join(

            cborutil.streamencode(

                {

                    b'name': b'command1',

                    b'args': {

                        b'foo': b'bar',

                        b'key1': b'val',

                    },

                }

            )

        )

        cbor3 = b''.join(

            cborutil.streamencode(

                {

                    b'name': b'command3',

                    b'args': {

                        b'biz': b'baz',

                        b'key': b'val',

                    },

                }

            )

        )

        results = list(