From afff154da4dd9eb9870e895cb143b2a2f67f19a2 Mon Sep 17 00:00:00 2001
From: tsukumi <tsukumijima@users.noreply.github.com>
Date: Sun, 10 Mar 2024 04:18:39 +0000
Subject: [PATCH] Add: test code for style-bert-vits2 as a library By executing
 "hatch run test:test", you can check whether the test passes in all Python
 3.9 to 3.12 environments.

---
 pyproject.toml                | 10 ++++++---
 style_bert_vits2/tts_model.py | 12 ++++++++---
 tests/.gitignore              |  1 +
 tests/test_main.py            | 39 +++++++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 6 deletions(-)
 create mode 100644 tests/.gitignore
 create mode 100644 tests/test_main.py

diff --git a/pyproject.toml b/pyproject.toml
index 4db5fe9..5995fc9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,24 +52,28 @@ Source = "https://github.com/litagin02/Style-Bert-VITS2"
 [tool.hatch.version]
 path = "style_bert_vits2/constants.py"
 
-[tool.hatch.envs.default]
+[tool.hatch.envs.test]
 dependencies = [
     "coverage[toml]>=6.5",
     "pytest",
 ]
-[tool.hatch.envs.default.scripts]
+[tool.hatch.envs.test.scripts]
+# Usage: `hatch run test:test`
 test = "pytest {args:tests}"
+# Usage: `hatch run test:coverage`
 test-cov = "coverage run -m pytest {args:tests}"
+# Usage: `hatch run test:cov-report`
 cov-report = [
     "- coverage combine",
     "coverage report",
 ]
+# Usage: `hatch run test:cov`
 cov = [
     "test-cov",
     "cov-report",
 ]
 
-[[tool.hatch.envs.all.matrix]]
+[[tool.hatch.envs.test.matrix]]
 python = ["3.9", "3.10", "3.11", "3.12"]
 
 [tool.coverage.run]
diff --git a/style_bert_vits2/tts_model.py b/style_bert_vits2/tts_model.py
index 109ef0e..fe4dcb5 100644
--- a/style_bert_vits2/tts_model.py
+++ b/style_bert_vits2/tts_model.py
@@ -1,6 +1,6 @@
 import warnings
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, TypedDict
 
 import gradio as gr
 import numpy as np
@@ -263,6 +263,13 @@ class TTSModel:
         return (self.hyper_parameters.data.sampling_rate, audio)
 
 
+class TTSModelInfo(TypedDict):
+    name: str
+    files: list[str]
+    styles: list[str]
+    speakers: list[str]
+
+
 class TTSModelHolder:
     """
     Style-Bert-Vits2 の音声合成モデルを管理するクラス。
@@ -297,8 +304,7 @@ class TTSModelHolder:
         self.model_files_dict: dict[str, list[Path]] = {}
         self.current_model: Optional[TTSModel] = None
         self.model_names: list[str] = []
-        self.models: list[TTSModel] = []
-        self.models_info: list[dict[str, Union[str, list[str]]]] = []
+        self.models_info: list[TTSModelInfo] = []
         self.refresh()
 
 
diff --git a/tests/.gitignore b/tests/.gitignore
new file mode 100644
index 0000000..697e56f
--- /dev/null
+++ b/tests/.gitignore
@@ -0,0 +1 @@
+*.wav
\ No newline at end of file
diff --git a/tests/test_main.py b/tests/test_main.py
new file mode 100644
index 0000000..597c232
--- /dev/null
+++ b/tests/test_main.py
@@ -0,0 +1,39 @@
+import pytest
+from scipy.io import wavfile
+
+from style_bert_vits2.constants import BASE_DIR
+from style_bert_vits2.tts_model import TTSModelHolder
+
+
+def synthesize(device: str = 'cpu'):
+
+    # モデル一覧を取得
+    model_holder = TTSModelHolder(BASE_DIR / 'model_assets', device)
+
+    # モデルが存在する場合、音声合成を実行
+    if len(model_holder.models_info) > 0:
+
+        # jvnv-F1-jp モデルを探す
+        for model_info in model_holder.models_info:
+            if model_info['name'] == 'jvnv-F1-jp':
+
+                # 音声合成を実行
+                model = model_holder.get_model(model_info['name'], model_info['files'][0])
+                model.load()
+                sample_rate, audio_data = model.infer("あらゆる現実を、すべて自分のほうへねじ曲げたのだ。")
+
+                # 音声データを保存
+                with open(BASE_DIR / 'tests/test.wav', mode='wb') as f:
+                    wavfile.write(f, sample_rate, audio_data)
+    else:
+        pytest.skip("音声合成モデルが見つかりませんでした。")
+
+
+def test_synthesize_cpu():
+    synthesize(device='cpu')
+    assert (BASE_DIR / 'tests/test.wav').exists()
+
+
+def test_synthesize_cuda():
+    synthesize(device='cuda')
+    assert (BASE_DIR / 'tests/test.wav').exists()