t/cli/test_sse.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from threading import Thread
import asyncio
from aiohttp import web
from aiohttp_sse import sse_response
from aiohttp_sse import EventSourceResponse
from datetime import datetime
import time
import os
import sys
import ssl


test_ssl = len(sys.argv) == 2 and sys.argv[1] == "ssl"


class SseResponse(EventSourceResponse):
    def __init__(self, **args):
        super().__init__(**args)
        self.headers.pop("X-Accel-Buffering", None)


async def events(request):
    async with sse_response(request, response_cls=SseResponse) as resp:
        for i in range(30000):
            data = "Server Time : {}".format(datetime.now())
            print(data)
            await resp.send(data)
            await asyncio.sleep(1)


def run_server():
    if test_ssl:
        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        ssl_context.load_cert_chain("t/certs/sse_server.crt", "t/certs/sse_server.key")
    app = web.Application()
    app.router.add_route("GET", "/events", events)
    if test_ssl:
        web.run_app(app, host="127.0.0.1", port=8080, ssl_context=ssl_context)
    else:
        web.run_app(app, host="127.0.0.1", port=8080)


def run_client():
    time.sleep(2)
    print("start testing sse...")

    import requests

    if test_ssl:
        url = "https://localhost:9443/events"
    else:
        url = "http://localhost:9080/events"
    headers = {"Accept": "text/event-stream"}
    response = requests.get(url, stream=True, headers=headers, verify=False)
    i = 0
    for line in response.iter_lines():
        if line:
            decoded = line.decode("utf-8")
            if decoded.startswith("data:"):
                print(decoded)
                i += 1
                if i == 3:
                    print("sse proxy test ok")
                    os._exit(0)


t = Thread(target=run_client, daemon=True)
t.start()
run_server()