tests/test-wireproto.py
changeset 36555 c28de190e20a
parent 36074 2f7290555c96
child 36801 66de4555cefd
equal deleted inserted replaced
36554:06f59bdde626 36555:c28de190e20a
     1 from __future__ import absolute_import, print_function
     1 from __future__ import absolute_import, print_function
     2 
     2 
     3 from mercurial import (
     3 from mercurial import (
     4     error,
     4     error,
       
     5     pycompat,
     5     util,
     6     util,
     6     wireproto,
     7     wireproto,
     7     wireprototypes,
     8     wireprototypes,
     8 )
     9 )
     9 stringio = util.stringio
    10 stringio = util.stringio
    11 class proto(object):
    12 class proto(object):
    12     def __init__(self, args):
    13     def __init__(self, args):
    13         self.args = args
    14         self.args = args
    14     def getargs(self, spec):
    15     def getargs(self, spec):
    15         args = self.args
    16         args = self.args
    16         args.setdefault('*', {})
    17         args.setdefault(b'*', {})
    17         names = spec.split()
    18         names = spec.split()
    18         return [args[n] for n in names]
    19         return [args[n] for n in names]
    19 
    20 
    20 class clientpeer(wireproto.wirepeer):
    21 class clientpeer(wireproto.wirepeer):
    21     def __init__(self, serverrepo):
    22     def __init__(self, serverrepo):
    24     @property
    25     @property
    25     def ui(self):
    26     def ui(self):
    26         return self.serverrepo.ui
    27         return self.serverrepo.ui
    27 
    28 
    28     def url(self):
    29     def url(self):
    29         return 'test'
    30         return b'test'
    30 
    31 
    31     def local(self):
    32     def local(self):
    32         return None
    33         return None
    33 
    34 
    34     def peer(self):
    35     def peer(self):
    39 
    40 
    40     def close(self):
    41     def close(self):
    41         pass
    42         pass
    42 
    43 
    43     def capabilities(self):
    44     def capabilities(self):
    44         return ['batch']
    45         return [b'batch']
    45 
    46 
    46     def _call(self, cmd, **args):
    47     def _call(self, cmd, **args):
       
    48         args = pycompat.byteskwargs(args)
    47         res = wireproto.dispatch(self.serverrepo, proto(args), cmd)
    49         res = wireproto.dispatch(self.serverrepo, proto(args), cmd)
    48         if isinstance(res, wireprototypes.bytesresponse):
    50         if isinstance(res, wireprototypes.bytesresponse):
    49             return res.data
    51             return res.data
    50         elif isinstance(res, bytes):
    52         elif isinstance(res, bytes):
    51             return res
    53             return res
    56         return stringio(self._call(cmd, **args))
    58         return stringio(self._call(cmd, **args))
    57 
    59 
    58     @wireproto.batchable
    60     @wireproto.batchable
    59     def greet(self, name):
    61     def greet(self, name):
    60         f = wireproto.future()
    62         f = wireproto.future()
    61         yield {'name': mangle(name)}, f
    63         yield {b'name': mangle(name)}, f
    62         yield unmangle(f.value)
    64         yield unmangle(f.value)
    63 
    65 
    64 class serverrepo(object):
    66 class serverrepo(object):
    65     def greet(self, name):
    67     def greet(self, name):
    66         return "Hello, " + name
    68         return b"Hello, " + name
    67 
    69 
    68     def filtered(self, name):
    70     def filtered(self, name):
    69         return self
    71         return self
    70 
    72 
    71 def mangle(s):
    73 def mangle(s):
    72     return ''.join(chr(ord(c) + 1) for c in s)
    74     return b''.join(pycompat.bytechr(ord(c) + 1) for c in pycompat.bytestr(s))
    73 def unmangle(s):
    75 def unmangle(s):
    74     return ''.join(chr(ord(c) - 1) for c in s)
    76     return b''.join(pycompat.bytechr(ord(c) - 1) for c in pycompat.bytestr(s))
    75 
    77 
    76 def greet(repo, proto, name):
    78 def greet(repo, proto, name):
    77     return mangle(repo.greet(unmangle(name)))
    79     return mangle(repo.greet(unmangle(name)))
    78 
    80 
    79 wireproto.commands['greet'] = (greet, 'name',)
    81 wireproto.commands[b'greet'] = (greet, b'name',)
    80 
    82 
    81 srv = serverrepo()
    83 srv = serverrepo()
    82 clt = clientpeer(srv)
    84 clt = clientpeer(srv)
    83 
    85 
    84 print(clt.greet("Foobar"))
    86 print(clt.greet(b"Foobar"))
    85 b = clt.iterbatch()
    87 b = clt.iterbatch()
    86 map(b.greet, ('Fo, =;:<o', 'Bar'))
    88 list(map(b.greet, (b'Fo, =;:<o', b'Bar')))
    87 b.submit()
    89 b.submit()
    88 print([r for r in b.results()])
    90 print([r for r in b.results()])