From 11beba1c4773e9af30566fda67aa856879bb30c7 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Sun, 5 Apr 2026 09:03:34 +0200
Subject: [PATCH] fix: clean up omnivoice import guard and __init__ error
 masking

Remove OmniVoice = None fallback in nodes/loader.py so missing omnivoice
gives a clear ImportError instead of a confusing AttributeError. Restore
__init__.py to clean form without the try/except that silently swallowed
real import errors. Add omnivoice mock to conftest.py and register a
pytest plugin that prevents pytest from treating the project root as a
Package node (which would try to import __init__.py outside a package
context and fail on the relative import).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 __init__.py       | 23 +++++++++--------------
 conftest.py       | 27 ++++++++++++++++++++++++++-
 nodes/loader.py   |  5 +----
 tests/conftest.py |  5 +++++
 4 files changed, 41 insertions(+), 19 deletions(-)
 create mode 100644 tests/conftest.py

diff --git a/__init__.py b/__init__.py
index c086790..571d1d5 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,18 +1,13 @@
-try:
-    from .nodes import OmniVoiceModelLoader, OmniVoiceGenerate
+from .nodes import OmniVoiceModelLoader, OmniVoiceGenerate
 
-    NODE_CLASS_MAPPINGS = {
-        "OmniVoiceModelLoader": OmniVoiceModelLoader,
-        "OmniVoiceGenerate": OmniVoiceGenerate,
-    }
+NODE_CLASS_MAPPINGS = {
+    "OmniVoiceModelLoader": OmniVoiceModelLoader,
+    "OmniVoiceGenerate": OmniVoiceGenerate,
+}
 
-    NODE_DISPLAY_NAME_MAPPINGS = {
-        "OmniVoiceModelLoader": "OmniVoice Model Loader",
-        "OmniVoiceGenerate": "OmniVoice Generate",
-    }
-except ImportError:
-    # Graceful fallback when loaded outside of a package context (e.g. pytest)
-    NODE_CLASS_MAPPINGS = {}
-    NODE_DISPLAY_NAME_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "OmniVoiceModelLoader": "OmniVoice Model Loader",
+    "OmniVoiceGenerate": "OmniVoice Generate",
+}
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
diff --git a/conftest.py b/conftest.py
index 0edeb20..fd668a6 100644
--- a/conftest.py
+++ b/conftest.py
@@ -1,6 +1,31 @@
 import sys
 import os
+from pathlib import Path
+from unittest.mock import MagicMock
 
 # Ensure the project root is on sys.path so `nodes.loader` can be imported
 # as a top-level package without needing to install the package.
-sys.path.insert(0, os.path.dirname(__file__))
+_root = os.path.dirname(__file__)
+sys.path.insert(0, _root)
+
+# Mock omnivoice so it can be imported in tests without being installed
+sys.modules.setdefault("omnivoice", MagicMock())
+
+# Prevent pytest from treating the project root as a Python Package node.
+# The root __init__.py is a ComfyUI extension entrypoint that uses relative
+# imports and cannot be imported by pytest without a proper package context.
+# Overriding pytest_collect_directory stops pytest from creating a Package
+# node (and calling Package.setup() which imports __init__.py) for the root.
+class _RootDirPlugin:
+    """Registered as a plugin so it intercepts hooks before conftest scope limits."""
+
+    def pytest_collect_directory(self, path, parent):
+        from _pytest.main import Dir
+        if Path(path) == Path(_root):
+            # Return a plain Dir node instead of a Package node for the root,
+            # so pytest skips importing __init__.py during Package.setup().
+            return Dir.from_parent(parent, path=Path(path))
+
+
+def pytest_configure(config):
+    config.pluginmanager.register(_RootDirPlugin(), "_root_dir_plugin")
diff --git a/nodes/loader.py b/nodes/loader.py
index b59c263..9686d3b 100644
--- a/nodes/loader.py
+++ b/nodes/loader.py
@@ -1,10 +1,7 @@
 import os
 import torch
 
-try:
-    from omnivoice import OmniVoice
-except ImportError:
-    OmniVoice = None  # type: ignore[assignment,misc]
+from omnivoice import OmniVoice
 
 try:
     import folder_paths
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..6f4a099
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,5 @@
+import sys
+from unittest.mock import MagicMock
+
+# Mock omnivoice so it can be imported in tests without being installed
+sys.modules.setdefault("omnivoice", MagicMock())