102 lines
2.6 KiB
Python
102 lines
2.6 KiB
Python
"""Test client that connects and sends infinite data."""
|
|
|
|
from __future__ import print_function
|
|
import argparse
|
|
import sys
|
|
|
|
from trollius import *
|
|
from trollius import test_utils
|
|
|
|
|
|
ARGS = argparse.ArgumentParser(description="TCP data sink example.")
|
|
ARGS.add_argument(
|
|
'--tls', action='store_true', dest='tls',
|
|
default=False, help='Use TLS')
|
|
ARGS.add_argument(
|
|
'--iocp', action='store_true', dest='iocp',
|
|
default=False, help='Use IOCP event loop (Windows only)')
|
|
ARGS.add_argument(
|
|
'--stop', action='store_true', dest='stop',
|
|
default=False, help='Stop the server by sending it b"stop" as data')
|
|
ARGS.add_argument(
|
|
'--host', action='store', dest='host',
|
|
default='127.0.0.1', help='Host name')
|
|
ARGS.add_argument(
|
|
'--port', action='store', dest='port',
|
|
default=1111, type=int, help='Port number')
|
|
ARGS.add_argument(
|
|
'--size', action='store', dest='size',
|
|
default=16*1024, type=int, help='Data size')
|
|
|
|
args = None
|
|
|
|
|
|
def dprint(*args):
|
|
print('source:', *args, file=sys.stderr)
|
|
|
|
|
|
class Client(Protocol):
|
|
|
|
total = 0
|
|
|
|
def connection_made(self, tr):
|
|
dprint('connecting to', tr.get_extra_info('peername'))
|
|
dprint('my socket is', tr.get_extra_info('sockname'))
|
|
self.tr = tr
|
|
self.lost = False
|
|
self.loop = get_event_loop()
|
|
self.waiter = Future()
|
|
if args.stop:
|
|
self.tr.write(b'stop')
|
|
self.tr.close()
|
|
else:
|
|
self.data = b'x'*args.size
|
|
self.write_some_data()
|
|
|
|
def write_some_data(self):
|
|
if self.lost:
|
|
dprint('lost already')
|
|
return
|
|
data = self.data
|
|
size = len(data)
|
|
self.total += size
|
|
dprint('writing', size, 'bytes; total', self.total)
|
|
self.tr.write(data)
|
|
self.loop.call_soon(self.write_some_data)
|
|
|
|
def connection_lost(self, exc):
|
|
dprint('lost connection', repr(exc))
|
|
self.lost = True
|
|
self.waiter.set_result(None)
|
|
|
|
|
|
@coroutine
|
|
def start(loop, host, port):
|
|
sslctx = None
|
|
if args.tls:
|
|
sslctx = test_utils.dummy_ssl_context()
|
|
tr, pr = yield From(loop.create_connection(Client, host, port,
|
|
ssl=sslctx))
|
|
dprint('tr =', tr)
|
|
dprint('pr =', pr)
|
|
yield From(pr.waiter)
|
|
|
|
|
|
def main():
|
|
global args
|
|
args = ARGS.parse_args()
|
|
if args.iocp:
|
|
from trollius.windows_events import ProactorEventLoop
|
|
loop = ProactorEventLoop()
|
|
set_event_loop(loop)
|
|
else:
|
|
loop = get_event_loop()
|
|
try:
|
|
loop.run_until_complete(start(loop, args.host, args.port))
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|