mirror of
https://github.com/ApfelTeeSaft/restream_playout.git
synced 2026-08-26 19:33:32 +00:00
81 lines
3.2 KiB
Python
81 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import unittest
|
|
|
|
from restream_playout.flv import amf0_encode
|
|
from restream_playout.rtmp import RTMP_COMMAND_AMF0, RTMP_DATA_AMF0, RTMPClient, amf0_decode_all
|
|
|
|
|
|
def command(*values: object) -> bytes:
|
|
return b"".join(amf0_encode(value) for value in values)
|
|
|
|
|
|
class RTMPTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_connect_publish_and_send(self) -> None:
|
|
media_received = asyncio.Event()
|
|
errors: list[BaseException] = []
|
|
|
|
async def server_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
peer = RTMPClient("rtmp://localhost/live/key")
|
|
try:
|
|
handshake = await reader.readexactly(1537)
|
|
c1 = handshake[1:]
|
|
s1 = b"\x00" * 1536
|
|
writer.write(b"\x03" + s1 + c1)
|
|
await writer.drain()
|
|
await reader.readexactly(1536)
|
|
peer._reader = reader
|
|
peer._writer = writer
|
|
connect_message = await peer.read_message()
|
|
self.assertEqual(amf0_decode_all(connect_message.payload)[0], "connect")
|
|
await peer.send_message(
|
|
RTMP_COMMAND_AMF0,
|
|
0,
|
|
command("_result", 1, {"fmsVer": "test"}, {"code": "NetConnection.Connect.Success"}),
|
|
stream_id=0,
|
|
chunk_stream_id=3,
|
|
)
|
|
create_message = await peer.read_message()
|
|
self.assertEqual(amf0_decode_all(create_message.payload)[0], "createStream")
|
|
await peer.send_message(
|
|
RTMP_COMMAND_AMF0,
|
|
0,
|
|
command("_result", 2, None, 1),
|
|
stream_id=0,
|
|
chunk_stream_id=3,
|
|
)
|
|
publish_message = await peer.read_message()
|
|
self.assertEqual(amf0_decode_all(publish_message.payload)[0], "publish")
|
|
await peer.send_message(
|
|
RTMP_COMMAND_AMF0,
|
|
0,
|
|
command("onStatus", 0, None, {"level": "status", "code": "NetStream.Publish.Start"}),
|
|
stream_id=1,
|
|
chunk_stream_id=3,
|
|
)
|
|
message = await peer.read_message()
|
|
self.assertEqual(message.message_type, RTMP_DATA_AMF0)
|
|
self.assertEqual(message.payload, b"hello")
|
|
media_received.set()
|
|
except BaseException as exc:
|
|
errors.append(exc)
|
|
media_received.set()
|
|
finally:
|
|
await peer.close()
|
|
|
|
server = await asyncio.start_server(server_handler, "127.0.0.1", 0)
|
|
port = server.sockets[0].getsockname()[1]
|
|
client = RTMPClient(f"rtmp://127.0.0.1:{port}/live/key")
|
|
try:
|
|
await client.connect()
|
|
await client.send_message(RTMP_DATA_AMF0, 0, b"hello")
|
|
await asyncio.wait_for(media_received.wait(), 2)
|
|
if errors:
|
|
raise errors[0]
|
|
finally:
|
|
await client.close()
|
|
server.close()
|
|
await server.wait_closed()
|
|
|