From 98adce8a2b8486da705f9be7f5b1eaf0926546f5 Mon Sep 17 00:00:00 2001 From: Henry Ruhs Date: Fri, 15 May 2026 11:46:51 +0200 Subject: [PATCH] Refactor RTC structure (#1113) * refactor rtc part1 * skip for macos * merge create spd and create sdp offer * fix lint * add test for create_sdp_offer * better naming for negotiate method as we get an answer * extend tests based on mutations * remove dead code * rename rtc store and related methods * clean store, move sender logic to stream helper under apis * generate tests for rtc store --- facefusion/apis/endpoints/stream.py | 6 +- facefusion/apis/stream_helper.py | 59 ++++++- facefusion/rtc.py | 235 +++++++++++++--------------- facefusion/rtc_store.py | 66 ++------ facefusion/types.py | 2 +- tests/test_api_stream.py | 9 +- tests/test_rtc.py | 71 ++++++--- tests/test_rtc_store.py | 57 ++++++- 8 files changed, 281 insertions(+), 224 deletions(-) diff --git a/facefusion/apis/endpoints/stream.py b/facefusion/apis/endpoints/stream.py index d4b2e95d..d0cc2bdb 100644 --- a/facefusion/apis/endpoints/stream.py +++ b/facefusion/apis/endpoints/stream.py @@ -3,9 +3,9 @@ from starlette.responses import Response from starlette.status import HTTP_201_CREATED, HTTP_404_NOT_FOUND from starlette.websockets import WebSocket -from facefusion import rtc_store, session_context, session_manager +from facefusion import session_context, session_manager from facefusion.apis.session_helper import extract_access_token -from facefusion.apis.stream_helper import handle_image_stream, handle_video_stream +from facefusion.apis.stream_helper import add_rtc_viewer, handle_image_stream, handle_video_stream # TODO: reject websocket with invalid or missing stream mode @@ -27,7 +27,7 @@ async def post_stream(request : Request) -> Response: if session_id: sdp_offer = await request.body() - sdp_answer = rtc_store.add_rtc_viewer(session_id, sdp_offer.decode()) + sdp_answer = add_rtc_viewer(session_id, sdp_offer.decode()) if sdp_answer: return Response(sdp_answer, status_code = HTTP_201_CREATED, media_type = 'application/sdp') diff --git a/facefusion/apis/stream_helper.py b/facefusion/apis/stream_helper.py index 1589ce41..be9876bb 100644 --- a/facefusion/apis/stream_helper.py +++ b/facefusion/apis/stream_helper.py @@ -1,20 +1,20 @@ import asyncio from collections import deque from collections.abc import AsyncIterator -from typing import Tuple, cast, get_args +from typing import Optional, Tuple, cast, get_args import cv2 import numpy from starlette.websockets import WebSocket, WebSocketState -from facefusion import rtc_store, session_context, session_manager, state_manager +from facefusion import rtc, rtc_store, session_context, session_manager, state_manager from facefusion.apis.api_helper import get_sec_websocket_protocol from facefusion.apis.session_helper import extract_access_token from facefusion.codecs.aom import create_aom_encoder, destroy_aom_encoder, encode_aom_buffer from facefusion.codecs.opus import create_opus_encoder, destroy_opus_encoder, encode_opus_buffer from facefusion.codecs.vpx import create_vpx_encoder, destroy_vpx_encoder, encode_vpx_buffer from facefusion.streamer import process_vision_frame -from facefusion.types import Resolution, SessionId, VideoCodec, VisionFrame +from facefusion.types import AudioCodec, PeerConnection, Resolution, RtcAudioTrack, RtcPeer, RtcVideoTrack, SdpAnswer, SdpOffer, SessionId, VideoCodec, VisionFrame async def receive_stream_frames(websocket : WebSocket) -> AsyncIterator[Tuple[int, bytes]]: @@ -42,6 +42,40 @@ async def receive_vision_frames(websocket : WebSocket) -> AsyncIterator[VisionFr websocket_event = await websocket.receive() +# TODO: clean up peer connection on failed sdp negotiation, wrap in run_in_executor to avoid blocking async event loop +def add_rtc_viewer(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpAnswer]: + rtc_peers = rtc_store.get_rtc_peers(session_id) + + if rtc_peers is not None: + payload_types = rtc.parse_sdp_payload_types(sdp_offer) + peer_connection : PeerConnection = rtc.create_peer_connection() + audio_codec : AudioCodec = 'opus' + audio_track : RtcAudioTrack = rtc.add_audio_track(peer_connection, 'sendonly', audio_codec, payload_types.get(audio_codec, 111)) + + #TODO: Fix me via resolve method + video_codec : VideoCodec = 'av1' + if payload_types.get('av1'): + video_codec = 'av1' + if payload_types.get('vp8'): + video_codec = 'vp8' + + video_track : RtcVideoTrack = rtc.add_video_track(peer_connection, 'sendonly', video_codec, payload_types.get(video_codec, 96)) + local_sdp = rtc.negotiate_sdp_answer(peer_connection, sdp_offer) + + if local_sdp: + rtc_peer : RtcPeer =\ + { + 'peer_connection': peer_connection, + 'video_track': video_track, + 'audio_track': audio_track + } + rtc_peers.append(rtc_peer) + + return local_sdp + + return None + + def run_aom_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : SessionId, initial_resolution : Resolution, keyframe_interval : int) -> None: aom_encoder = create_aom_encoder(initial_resolution, 4500, 8, 10) current_resolution = initial_resolution @@ -65,7 +99,10 @@ def run_aom_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : Se frame_buffer = encode_aom_buffer(aom_encoder, yuv_frame.tobytes(), frame_resolution, pts) if frame_buffer: - rtc_store.send_rtc_video(session_id, frame_buffer) + rtc_peers = rtc_store.get_rtc_peers(session_id) + + if rtc_peers: + rtc.send_video_to_peers(rtc_peers, frame_buffer) pts += 1 @@ -96,7 +133,10 @@ def run_vp8_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : Se frame_buffer = encode_vpx_buffer(vpx_encoder, yuv_frame.tobytes(), frame_resolution, pts) if frame_buffer: - rtc_store.send_rtc_video(session_id, frame_buffer) + rtc_peers = rtc_store.get_rtc_peers(session_id) + + if rtc_peers: + rtc.send_video_to_peers(rtc_peers, frame_buffer) pts += 1 @@ -159,7 +199,7 @@ async def handle_video_stream(websocket : WebSocket) -> None: audio_timestamp = 0 vision_frame_deque.append(first_vision_frame) - rtc_store.create_rtc_stream(session_id) + rtc_store.create_rtc_peers(session_id) event_loop = asyncio.get_running_loop() encode_loop = run_aom_encode_loop @@ -186,7 +226,10 @@ async def handle_video_stream(websocket : WebSocket) -> None: audio_buffer = encode_opus_buffer(opus_encoder, audio_chunk.tobytes(), 960) if audio_buffer: - rtc_store.send_rtc_audio(session_id, audio_buffer, audio_timestamp) + rtc_peers = rtc_store.get_rtc_peers(session_id) + + if rtc_peers: + rtc.send_audio_to_peers(rtc_peers, audio_buffer, audio_timestamp) audio_timestamp += 960 @@ -196,7 +239,7 @@ async def handle_video_stream(websocket : WebSocket) -> None: if opus_encoder: destroy_opus_encoder(opus_encoder) - rtc_store.destroy_rtc_stream(session_id) + rtc_store.destroy_rtc_peers(session_id) if websocket.client_state == WebSocketState.CONNECTED: await websocket.close() diff --git a/facefusion/rtc.py b/facefusion/rtc.py index a8f0b1f6..c8b66e11 100644 --- a/facefusion/rtc.py +++ b/facefusion/rtc.py @@ -43,115 +43,9 @@ def create_peer_connection( return datachannel_library.rtcCreatePeerConnection(ctypes.byref(rtc_configuration)) -#TODO: I think we dont need this method at all -def build_media_description(media_type : str, payload_type : int, rtp_codec : str, media_direction : MediaDirection, media_id : int) -> bytes: - lines =\ - [ - 'm=' + media_type + ' 9 UDP/TLS/RTP/SAVPF ' + str(payload_type), - 'a=rtpmap:' + str(payload_type) + ' ' + rtp_codec, - 'a=rtcp-fb:' + str(payload_type) + ' nack', - 'a=rtcp-fb:' + str(payload_type) + ' nack pli', - 'a=' + media_direction, - 'a=mid:' + str(media_id), - 'a=rtcp-mux', - '' - ] - return '\r\n'.join(lines).encode() - - -def parse_sdp_payload_types(sdp_offer : SdpOffer) -> Dict[str, int]: - payload_types : Dict[str, int] = {} - - # TODO: consider having a codec helper to resolve these - for line in sdp_offer.splitlines(): - if line.startswith('a=rtpmap:') and 'AV1/90000' in line and not payload_types.get('av1'): - payload_types['av1'] = int(line.split(':')[1].split(' ')[0]) - if line.startswith('a=rtpmap:') and 'VP8/90000' in line and not payload_types.get('vp8'): - payload_types['vp8'] = int(line.split(':')[1].split(' ')[0]) - if line.startswith('a=rtpmap:') and 'opus/48000/2' in line and not payload_types.get('opus'): - payload_types['opus'] = int(line.split(':')[1].split(' ')[0]) - - return payload_types - - -def add_audio_track(peer_connection : PeerConnection, media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> RtcAudioTrack: +# TODO: check if sleep is needed +def create_sdp_offer(peer_connection : PeerConnection) -> Optional[SdpOffer]: datachannel_library = datachannel_module.create_static_library() - - # TODO: Fix me via resolve method - rtp_codec = 'opus/48000/2' - if audio_codec == 'opus': - rtp_codec = 'opus/48000/2' - - media_description = build_media_description('audio', payload_type, rtp_codec, media_direction, 1) - - audio_track = datachannel_library.rtcAddTrack(peer_connection, media_description) - - audio_packetizer = datachannel_module.define_rtc_packetizer_init() - audio_packetizer.ssrc = 43 - audio_packetizer.cname = b'audio' - audio_packetizer.payloadType = payload_type - audio_packetizer.clockRate = 48000 - - if audio_codec == 'opus': - datachannel_library.rtcSetOpusPacketizer(audio_track, ctypes.byref(audio_packetizer)) - - datachannel_library.rtcChainRtcpSrReporter(audio_track) - - return audio_track - - -def add_video_track(peer_connection : PeerConnection, media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> RtcVideoTrack: - datachannel_library = datachannel_module.create_static_library() - - #TODO: Fix me via resolve method - rtp_codec = 'AV1/90000' - if video_codec == 'av1': - rtp_codec = 'AV1/90000' - if video_codec == 'vp8': - rtp_codec = 'VP8/90000' - - media_description = build_media_description('video', payload_type, rtp_codec, media_direction, 0) - - video_track = datachannel_library.rtcAddTrack(peer_connection, media_description) - - video_packetizer = datachannel_module.define_rtc_packetizer_init() - video_packetizer.ssrc = 42 - video_packetizer.cname = b'video' - video_packetizer.payloadType = payload_type - video_packetizer.clockRate = 90000 - video_packetizer.maxFragmentSize = 1200 - - if video_codec == 'av1': - video_packetizer.obuPacketization = 1 - datachannel_library.rtcSetAV1Packetizer(video_track, ctypes.byref(video_packetizer)) - if video_codec == 'vp8': - datachannel_library.rtcSetVP8Packetizer(video_track, ctypes.byref(video_packetizer)) - - datachannel_library.rtcChainRtcpSrReporter(video_track) - datachannel_library.rtcChainRtcpNackResponder(video_track, 512) - - return video_track - - -def create_sdp(peer_connection : PeerConnection) -> Optional[SdpOffer]: - datachannel_library = datachannel_module.create_static_library() - datachannel_library.rtcSetLocalDescription(peer_connection, b'offer') - buffer_size = 8192 - buffer_string = ctypes.create_string_buffer(buffer_size) - - if datachannel_library.rtcGetLocalDescription(peer_connection, buffer_string, buffer_size) > 0: - return buffer_string.value.decode() - - return None - - -# TODO: move from testing suite helper to rtc.py - belongs here to complete the rtc flow -def create_sdp_offer() -> Optional[SdpOffer]: - datachannel_library = datachannel_module.create_static_library() - peer_connection = create_peer_connection(disable_auto_negotiation = True) - - datachannel_library.rtcAddTrack(peer_connection, build_media_description('video', 96, 'VP8/90000', 'recvonly', 0)) - datachannel_library.rtcAddTrack(peer_connection, build_media_description('audio', 111, 'opus/48000/2', 'recvonly', 1)) datachannel_library.rtcSetLocalDescription(peer_connection, b'offer') buffer_size = 16384 @@ -160,31 +54,15 @@ def create_sdp_offer() -> Optional[SdpOffer]: while time.monotonic() < wait_limit: if datachannel_library.rtcGetLocalDescription(peer_connection, buffer_string, buffer_size) > 0: - sdp = buffer_string.value.decode() - datachannel_library.rtcDeletePeerConnection(peer_connection) - return sdp + return buffer_string.value.decode() time.sleep(0.05) - datachannel_library.rtcDeletePeerConnection(peer_connection) return None -@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_void_p) -def on_sdp_ready(peer_connection : int, sdp : Optional[bytes], sdp_type : int, user_pointer : Optional[int]) -> None: - ctypes.cast(user_pointer, ctypes.py_object).value.set() - - -# TODO: unused callback, remove or wire up -@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_int, ctypes.c_void_p) -def on_ice_complete(peer_connection : int, state : int, user_pointer : Optional[int]) -> None: - if state == 2: - context = ctypes.cast(user_pointer, ctypes.py_object).value - context['event'].set() - - # TODO: sanitize sdp_offer, wrap in run_in_executor, track peer connection state -def negotiate_sdp(peer_connection : PeerConnection, sdp_offer : SdpOffer) -> Optional[SdpAnswer]: +def negotiate_sdp_answer(peer_connection : PeerConnection, sdp_offer : SdpOffer) -> Optional[SdpAnswer]: datachannel_library = datachannel_module.create_static_library() sdp_event = threading.Event() sdp_event_pointer = ctypes.cast(id(sdp_event), ctypes.c_void_p) @@ -249,3 +127,108 @@ def delete_peers(rtc_peers : List[RtcPeer]) -> None: datachannel_library.rtcDeletePeerConnection(peer_connection_id) return None + + +def add_audio_track(peer_connection : PeerConnection, media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> RtcAudioTrack: + datachannel_library = datachannel_module.create_static_library() + media_description = create_audio_description(media_direction, audio_codec, payload_type) + + audio_track = datachannel_library.rtcAddTrack(peer_connection, media_description) + + audio_packetizer = datachannel_module.define_rtc_packetizer_init() + audio_packetizer.ssrc = 43 + audio_packetizer.cname = b'audio' + audio_packetizer.payloadType = payload_type + audio_packetizer.clockRate = 48000 + + if audio_codec == 'opus': + datachannel_library.rtcSetOpusPacketizer(audio_track, ctypes.byref(audio_packetizer)) + + datachannel_library.rtcChainRtcpSrReporter(audio_track) + + return audio_track + + +def add_video_track(peer_connection : PeerConnection, media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> RtcVideoTrack: + datachannel_library = datachannel_module.create_static_library() + media_description = create_video_description(media_direction, video_codec, payload_type) + + video_track = datachannel_library.rtcAddTrack(peer_connection, media_description) + + video_packetizer = datachannel_module.define_rtc_packetizer_init() + video_packetizer.ssrc = 42 + video_packetizer.cname = b'video' + video_packetizer.payloadType = payload_type + video_packetizer.clockRate = 90000 + video_packetizer.maxFragmentSize = 1200 + + if video_codec == 'av1': + video_packetizer.obuPacketization = 1 + datachannel_library.rtcSetAV1Packetizer(video_track, ctypes.byref(video_packetizer)) + if video_codec == 'vp8': + datachannel_library.rtcSetVP8Packetizer(video_track, ctypes.byref(video_packetizer)) + + datachannel_library.rtcChainRtcpSrReporter(video_track) + datachannel_library.rtcChainRtcpNackResponder(video_track, 512) + + return video_track + + +def create_audio_description(media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> bytes: + rtp_codec = 'opus/48000/2' + if audio_codec == 'opus': + rtp_codec = 'opus/48000/2' + + lines =\ + [ + 'm=audio 9 UDP/TLS/RTP/SAVPF ' + str(payload_type), + 'a=rtpmap:' + str(payload_type) + ' ' + rtp_codec, + 'a=rtcp-fb:' + str(payload_type) + ' nack', + 'a=rtcp-fb:' + str(payload_type) + ' nack pli', + 'a=' + media_direction, + 'a=mid:1', + 'a=rtcp-mux', + '' + ] + return '\r\n'.join(lines).encode() + + +def create_video_description(media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> bytes: + rtp_codec = 'AV1/90000' + if video_codec == 'av1': + rtp_codec = 'AV1/90000' + if video_codec == 'vp8': + rtp_codec = 'VP8/90000' + + lines =\ + [ + 'm=video 9 UDP/TLS/RTP/SAVPF ' + str(payload_type), + 'a=rtpmap:' + str(payload_type) + ' ' + rtp_codec, + 'a=rtcp-fb:' + str(payload_type) + ' nack', + 'a=rtcp-fb:' + str(payload_type) + ' nack pli', + 'a=' + media_direction, + 'a=mid:0', + 'a=rtcp-mux', + '' + ] + return '\r\n'.join(lines).encode() + + +def parse_sdp_payload_types(sdp_offer : SdpOffer) -> Dict[str, int]: + payload_types : Dict[str, int] = {} + + # TODO: consider having a codec helper to resolve these + for line in sdp_offer.splitlines(): + if line.startswith('a=rtpmap:') and 'AV1/90000' in line and not payload_types.get('av1'): + payload_types['av1'] = int(line.split(':')[1].split(' ')[0]) + if line.startswith('a=rtpmap:') and 'VP8/90000' in line and not payload_types.get('vp8'): + payload_types['vp8'] = int(line.split(':')[1].split(' ')[0]) + if line.startswith('a=rtpmap:') and 'opus/48000/2' in line and not payload_types.get('opus'): + payload_types['opus'] = int(line.split(':')[1].split(' ')[0]) + + return payload_types + + +@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_void_p) +def on_sdp_ready(peer_connection : int, sdp : Optional[bytes], sdp_type : int, user_pointer : Optional[int]) -> None: + ctypes.cast(user_pointer, ctypes.py_object).value.set() diff --git a/facefusion/rtc_store.py b/facefusion/rtc_store.py index 9725ad45..41970573 100644 --- a/facefusion/rtc_store.py +++ b/facefusion/rtc_store.py @@ -1,70 +1,22 @@ -from typing import List, Optional +from typing import List from facefusion import rtc -from facefusion.types import AudioCodec, PeerConnection, RtcAudioTrack, RtcPeer, RtcStreamStore, RtcVideoTrack, SdpAnswer, SdpOffer, SessionId, VideoCodec +from facefusion.types import RtcPeer, RtcStore, SessionId -# TODO: aint this a peer store? -RTC_STREAM_STORE : RtcStreamStore = {} +RTC_STORE : RtcStore = {} -def get_rtc_stream(session_id : SessionId) -> Optional[List[RtcPeer]]: - return RTC_STREAM_STORE.get(session_id) +def get_rtc_peers(session_id : SessionId) -> List[RtcPeer]: + return RTC_STORE.get(session_id) -def create_rtc_stream(session_id : SessionId) -> None: - RTC_STREAM_STORE[session_id] = [] +def create_rtc_peers(session_id : SessionId) -> None: + RTC_STORE[session_id] = [] -def destroy_rtc_stream(session_id : SessionId) -> None: - rtc_peers = RTC_STREAM_STORE.pop(session_id, None) +def destroy_rtc_peers(session_id : SessionId) -> None: + rtc_peers = RTC_STORE.pop(session_id, None) if rtc_peers: rtc.delete_peers(rtc_peers) - - -# TODO: clean up peer connection on failed sdp negotiation, wrap in run_in_executor to avoid blocking async event loop -def add_rtc_viewer(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpAnswer]: - if session_id in RTC_STREAM_STORE: - payload_types = rtc.parse_sdp_payload_types(sdp_offer) - peer_connection : PeerConnection = rtc.create_peer_connection() - audio_codec : AudioCodec = 'opus' - audio_track : RtcAudioTrack = rtc.add_audio_track(peer_connection, 'sendonly', audio_codec, payload_types.get(audio_codec, 111)) - - #TODO: Fix me via resolve method - video_codec : VideoCodec = 'av1' - if payload_types.get('av1'): - video_codec = 'av1' - if payload_types.get('vp8'): - video_codec = 'vp8' - - video_track : RtcVideoTrack = rtc.add_video_track(peer_connection, 'sendonly', video_codec, payload_types.get(video_codec, 96)) - local_sdp = rtc.negotiate_sdp(peer_connection, sdp_offer) - - if local_sdp: - rtc_peer : RtcPeer =\ - { - 'peer_connection': peer_connection, - 'video_track': video_track, - 'audio_track': audio_track - } - RTC_STREAM_STORE[session_id].append(rtc_peer) - - return local_sdp - - return None - - -# TODO: detect and remove dead peers -def send_rtc_video(session_id : SessionId, frame_buffer : bytes) -> None: - rtc_peers = get_rtc_stream(session_id) - - if rtc_peers: - rtc.send_video_to_peers(rtc_peers, frame_buffer) - - -def send_rtc_audio(session_id : SessionId, audio_buffer : bytes, audio_pts : int) -> None: - rtc_peers = get_rtc_stream(session_id) - - if rtc_peers: - rtc.send_audio_to_peers(rtc_peers, audio_buffer, audio_pts) diff --git a/facefusion/types.py b/facefusion/types.py index 7b2e17b3..9856e686 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -286,7 +286,7 @@ RtcPeer = TypedDict('RtcPeer', 'audio_track': RtcAudioTrack, }) -RtcStreamStore : TypeAlias = Dict[SessionId, List[RtcPeer]] +RtcStore : TypeAlias = Dict[SessionId, List[RtcPeer]] ModelOptions : TypeAlias = Dict[str, Any] ModelSet : TypeAlias = Dict[str, ModelOptions] diff --git a/tests/test_api_stream.py b/tests/test_api_stream.py index c3104a8e..a629b6f3 100644 --- a/tests/test_api_stream.py +++ b/tests/test_api_stream.py @@ -13,6 +13,7 @@ from facefusion.apis.core import create_api from facefusion.core import common_pre_check from facefusion.download import conditional_download from facefusion.hash_helper import create_hash +from facefusion.libraries import datachannel as datachannel_module from .assert_helper import get_test_example_file, get_test_examples_directory @@ -123,7 +124,7 @@ def test_stream_video(test_client : TestClient, create_event : threading.Event, 'Authorization': 'Bearer ' + access_token }) - with patch('facefusion.rtc_store.send_rtc_video', side_effect = partial(set_event, event = create_event)): + with patch('facefusion.rtc.send_video_to_peers', side_effect = partial(set_event, event = create_event)): with test_client.websocket_connect('/stream?mode=video&codec=' + video_codec, subprotocols = [ 'access_token.' + access_token @@ -131,7 +132,11 @@ def test_stream_video(test_client : TestClient, create_event : threading.Event, websocket.send_bytes(chr(1).encode() + source_content) websocket.receive_text() - sdp_offer = rtc.create_sdp_offer() + peer_connection = rtc.create_peer_connection(disable_auto_negotiation = True) + rtc.add_video_track(peer_connection, 'recvonly', 'vp8', 96) + rtc.add_audio_track(peer_connection, 'recvonly', 'opus', 111) + sdp_offer = rtc.create_sdp_offer(peer_connection) + datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection) stream_response = test_client.post('/stream', content = sdp_offer, headers = { 'Authorization': 'Bearer ' + access_token, diff --git a/tests/test_rtc.py b/tests/test_rtc.py index a0d221d0..88ebbd2c 100644 --- a/tests/test_rtc.py +++ b/tests/test_rtc.py @@ -16,12 +16,6 @@ def before_all() -> None: vpx_module.pre_check() -# TODO: add test_parse_sdp_payload_types -def test_build_media_description() -> None: - assert rtc.build_media_description('audio', 111, 'opus/48000/2', 'sendonly', 1) == b'm=audio 9 UDP/TLS/RTP/SAVPF 111\r\na=rtpmap:111 opus/48000/2\r\na=rtcp-fb:111 nack\r\na=rtcp-fb:111 nack pli\r\na=sendonly\r\na=mid:1\r\na=rtcp-mux\r\n' - assert rtc.build_media_description('video', 96, 'VP8/90000', 'recvonly', 0) == b'm=video 9 UDP/TLS/RTP/SAVPF 96\r\na=rtpmap:96 VP8/90000\r\na=rtcp-fb:96 nack\r\na=rtcp-fb:96 nack pli\r\na=recvonly\r\na=mid:0\r\na=rtcp-mux\r\n' - - def test_create_peer_connection() -> None: peer_connection = rtc.create_peer_connection() datachannel_library = datachannel_module.create_static_library() @@ -30,40 +24,41 @@ def test_create_peer_connection() -> None: assert datachannel_library.rtcDeletePeerConnection(peer_connection) == 0 -def test_add_audio_track() -> None: - peer_connection = rtc.create_peer_connection() +def test_create_sdp_offer() -> None: + peer_connection = rtc.create_peer_connection(disable_auto_negotiation = True) + rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96) + rtc.add_audio_track(peer_connection, 'sendonly', 'opus', 111) + sdp_offer = rtc.create_sdp_offer(peer_connection) - assert rtc.add_audio_track(peer_connection, 'sendonly', 'opus', 111) > 0 + assert sdp_offer + assert 'm=video' in sdp_offer + assert 'VP8/90000' in sdp_offer + assert 'm=audio' in sdp_offer + assert 'opus/48000/2' in sdp_offer datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection) -def test_add_video_track() -> None: - peer_connection = rtc.create_peer_connection() - - assert rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96) > 0 - - datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection) - - -def test_negotiate_sdp() -> None: +def test_negotiate_sdp_answer() -> None: datachannel_library = datachannel_module.create_static_library() sender_connection = rtc.create_peer_connection() rtc.add_video_track(sender_connection, 'sendonly', 'vp8', 96) rtc.add_audio_track(sender_connection, 'sendonly', 'opus', 111) - sdp_offer = rtc.create_sdp(sender_connection) + sdp_offer = rtc.create_sdp_offer(sender_connection) receiver_connection = rtc.create_peer_connection() rtc.add_video_track(receiver_connection, 'recvonly', 'vp8', 96) rtc.add_audio_track(receiver_connection, 'recvonly', 'opus', 111) - sdp_answer = rtc.negotiate_sdp(receiver_connection, sdp_offer) + sdp_answer = rtc.negotiate_sdp_answer(receiver_connection, sdp_offer) assert sdp_answer assert 'm=video' in sdp_answer assert 'VP8/90000' in sdp_answer assert 'm=audio' in sdp_answer assert 'opus/48000/2' in sdp_answer + # TODO: review + assert 'a=recvonly' in sdp_answer assert datachannel_library.rtcDeletePeerConnection(sender_connection) == 0 assert datachannel_library.rtcDeletePeerConnection(receiver_connection) == 0 @@ -84,3 +79,39 @@ def test_delete_peers() -> None: rtc.delete_peers(rtc_peers) assert datachannel_library.rtcDeletePeerConnection(peer_connection) < 0 + + +def test_add_audio_track() -> None: + peer_connection = rtc.create_peer_connection() + audio_track = rtc.add_audio_track(peer_connection, 'sendonly', 'opus', 111) + + assert audio_track > 0 + + # TODO: review + sdp_offer = rtc.create_sdp_offer(peer_connection) + + assert 'opus/48000/2' in sdp_offer + + datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection) + + +def test_add_video_track() -> None: + peer_connection = rtc.create_peer_connection() + video_track = rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96) + + assert video_track > 0 + + # TODO: review + sdp_offer = rtc.create_sdp_offer(peer_connection) + + assert 'VP8/90000' in sdp_offer + + datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection) + + +def test_create_audio_description() -> None: + assert rtc.create_audio_description('sendonly', 'opus', 111) == b'm=audio 9 UDP/TLS/RTP/SAVPF 111\r\na=rtpmap:111 opus/48000/2\r\na=rtcp-fb:111 nack\r\na=rtcp-fb:111 nack pli\r\na=sendonly\r\na=mid:1\r\na=rtcp-mux\r\n' + + +def test_create_video_description() -> None: + assert rtc.create_video_description('recvonly', 'vp8', 96) == b'm=video 9 UDP/TLS/RTP/SAVPF 96\r\na=rtpmap:96 VP8/90000\r\na=rtcp-fb:96 nack\r\na=rtcp-fb:96 nack pli\r\na=recvonly\r\na=mid:0\r\na=rtcp-mux\r\n' diff --git a/tests/test_rtc_store.py b/tests/test_rtc_store.py index de7ae6b8..2e8d6c40 100644 --- a/tests/test_rtc_store.py +++ b/tests/test_rtc_store.py @@ -1,7 +1,10 @@ +from typing import List + import pytest -from facefusion import state_manager +from facefusion import rtc, rtc_store, state_manager from facefusion.libraries import datachannel as datachannel_module, opus as opus_module, vpx as vpx_module +from facefusion.types import RtcPeer @pytest.fixture(scope = 'module', autouse = True) @@ -13,11 +16,51 @@ def before_all() -> None: vpx_module.pre_check() -# TODO: test create_rtc_stream, get_rtc_stream, destroy_rtc_stream lifecycle -def test_rtc_stream_lifecycle() -> None: - pass +@pytest.fixture(autouse = True) +def before_each() -> None: + rtc_store.RTC_STORE.clear() -# TODO: test add_rtc_viewer with valid session and sdp offer -def test_add_rtc_viewer() -> None: - pass +# TODO: needs review +def test_create_rtc_peers() -> None: + rtc_store.create_rtc_peers('test-session') + + assert rtc_store.RTC_STORE.get('test-session') == [] + + +# TODO: needs review +def test_get_rtc_peers() -> None: + assert rtc_store.get_rtc_peers('test-session') is None + + rtc_store.create_rtc_peers('test-session') + + assert rtc_store.get_rtc_peers('test-session') == [] + + +# TODO: needs review +def test_destroy_rtc_peers() -> None: + rtc_store.create_rtc_peers('test-session') + rtc_store.destroy_rtc_peers('test-session') + + assert rtc_store.get_rtc_peers('test-session') is None + + +# TODO: needs review +def test_destroy_rtc_peers_with_connections() -> None: + datachannel_library = datachannel_module.create_static_library() + peer_connection = rtc.create_peer_connection() + rtc_store.create_rtc_peers('test-session') + rtc_peers : List[RtcPeer] =\ + [ + { + 'peer_connection': peer_connection, + 'video_track': 0, + 'audio_track': 0 + } + ] + rtc_store.RTC_STORE['test-session'] = rtc_peers + + rtc_store.destroy_rtc_peers('test-session') + + assert rtc_store.get_rtc_peers('test-session') is None + assert datachannel_library.rtcDeletePeerConnection(peer_connection) < 0