From 873c86d0c89888de6c6fb75a5036cf4b1e66bff7 Mon Sep 17 00:00:00 2001
From: BangProx <stush0626@gmail.com>
Date: Fri, 16 Aug 2024 20:07:14 +0900
Subject: [PATCH] =?UTF-8?q?websocket=5Ftest.py=20=EC=B6=94=EA=B0=80?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 routers/websocket.py      |   4 +-
 routers/websocket_test.py | 101 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 2 deletions(-)
 create mode 100644 routers/websocket_test.py

diff --git a/routers/websocket.py b/routers/websocket.py
index bcf4061..099eb97 100644
--- a/routers/websocket.py
+++ b/routers/websocket.py
@@ -65,8 +65,8 @@ async def websocket_task(websocket: WebSocket):
     print("[WebSocketTask]-[****] Model Loaded Successfully with", device)
 
     # Accept WebSocket
-    connection_uuid = await websocket.receive_bytes()
-    print("[WebSocketTask] Connection [" + connection_uuid + "] Accepted")
+    #connection_uuid = await websocket.receive_bytes()
+    #print("[WebSocketTask] Connection [" + connection_uuid + "] Accepted")
 
     # Execute TTS Thread until WebSocket Disconnected
     ttsThread = threading.Thread(target=textToSpeech, args=stop_event)
diff --git a/routers/websocket_test.py b/routers/websocket_test.py
new file mode 100644
index 0000000..0cacabe
--- /dev/null
+++ b/routers/websocket_test.py
@@ -0,0 +1,101 @@
+from fastapi import APIRouter, WebSocket
+import asyncio
+import uuid
+import json
+from queue import Queue
+from datetime import datetime, timedelta
+from time import sleep
+import threading
+import wave # <- 임시로 import하는거. 추후 삭제 필요.
+
+router = APIRouter()
+
+async def read_wav(file_path):
+    with wave.open(file_path, 'rb') as wav_file:
+        frames = wav_file.readframes(wav_file.getnframes())
+    return frames
+
+async def textToSpeech(stop_event):
+    print("[TTSThread] TTS Thread Executed")
+
+    while not stop_event.is_set():
+        sleep(3)
+        try:
+            now = datetime.utcnow()
+            # pop text from the queue.
+            if not text_queue.empty():
+                processed_result = read_wav("../received_audio.wav")
+                if processed_result != None:
+                    result_queue.put(processed_result)
+
+        except Exception as e:
+            print(f"[TTSThread] Error processing text data: {e}")
+            break
+
+    print("[TTSThread] TTS Thread Destroyed")
+
+async def process_thread(websocket: WebSocket):
+    print("[ProcessTask] Process Task Initiated")
+    while not stop_event.is_set():
+        await asyncio.sleep(0.25)
+        try:
+            if not result_queue.empty():
+                audio_result = result_queue.get()
+                #audio_id = str(uuid.uuid4())
+                await websocket.send(audio_result)
+                
+        except Exception as e:
+            print(f"[ProcessTask] Error : {e}")
+            break
+    
+    print("[ProcessTask] Process Task Destroyed")
+
+
+
+#======= WebSocket ========#
+@router.websocket("/tts/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    print("Configuring BE Socket")
+    await websocket.accept()
+    print("BE Socket Accepted")
+    try:
+        while True:
+            data = await websocket.receive_text()
+            print("data:\n",data)
+            processed_result = read_wav("../received_audio.wav")
+            await websocket.send_bytes(processed_result)
+    except Exception as e:
+        print("Exception as {e}")
+    finally:
+        print("WebSocket connection closed.")
+
+
+
+@router.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    print("Configuring BE Socket")
+    await websocket.accept()
+    print("BE Socket Accepted")
+    
+    while True:
+        data = await websocket.receive_text()
+        await websocket.send_text(f"Echo: {data}")
+
+@router.websocket("/audio/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    print("WebSocket connection accepted.")
+    
+    try:
+        while True:
+            # 음성 데이터를 바이너리로 수신
+            data = await websocket.receive_bytes()
+            # 데이터 처리 예: 파일로 저장
+            with open("received_audio.wav", "wb") as f:
+                f.write(data)
+            # 클라이언트에 데이터 전송 (예: 수신 완료 알림)
+            await websocket.send_text("Audio received and saved.")
+    except Exception as e:
+        print("Error:", e)
+    finally:
+        print("WebSocket connection closed.")