139 lines
4.5 KiB
Python
139 lines
4.5 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
|
|
class FakeRequestException(Exception):
|
|
def __init__(self, message="", response=None):
|
|
super().__init__(message)
|
|
self.response = response
|
|
|
|
|
|
class FakeTimeout(FakeRequestException):
|
|
pass
|
|
|
|
|
|
def install_dependency_stubs():
|
|
requests_module = types.ModuleType("requests")
|
|
exceptions = types.SimpleNamespace(
|
|
RequestException=FakeRequestException,
|
|
Timeout=FakeTimeout,
|
|
)
|
|
requests_module.exceptions = exceptions
|
|
requests_module.get = Mock()
|
|
requests_module.post = Mock()
|
|
sys.modules["requests"] = requests_module
|
|
|
|
torch_module = types.ModuleType("torch")
|
|
torch_module.float32 = "float32"
|
|
torch_module.Tensor = type("Tensor", (), {})
|
|
torch_module.zeros = Mock(return_value="placeholder-image")
|
|
sys.modules["torch"] = torch_module
|
|
|
|
tiktoken_module = types.ModuleType("tiktoken")
|
|
tiktoken_module.get_encoding = Mock()
|
|
sys.modules["tiktoken"] = tiktoken_module
|
|
|
|
pil_module = types.ModuleType("PIL")
|
|
image_module = types.ModuleType("PIL.Image")
|
|
image_module.open = Mock()
|
|
image_module.fromarray = Mock()
|
|
pil_module.Image = image_module
|
|
sys.modules["PIL"] = pil_module
|
|
sys.modules["PIL.Image"] = image_module
|
|
|
|
sys.modules.setdefault("numpy", types.ModuleType("numpy"))
|
|
|
|
|
|
def load_node_module():
|
|
install_dependency_stubs()
|
|
root = Path(__file__).resolve().parents[1]
|
|
package_name = "comfy_openrouter_node_test"
|
|
|
|
package = types.ModuleType(package_name)
|
|
package.__path__ = [str(root)]
|
|
sys.modules[package_name] = package
|
|
|
|
spec = importlib.util.spec_from_file_location(
|
|
f"{package_name}.node",
|
|
root / "node.py",
|
|
submodule_search_locations=[str(root)],
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class RequestTimeoutTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.node_module = load_node_module()
|
|
self.node = self.node_module.OpenRouterNode()
|
|
self.node.fetch_credits = Mock(return_value="Remaining: $1.000")
|
|
self.node.count_tokens = Mock(return_value=1)
|
|
|
|
def call_generate_response(self, request_timeout=120):
|
|
return self.node.generate_response(
|
|
api_key="test-key",
|
|
system_prompt="system",
|
|
user_message_box="hello",
|
|
model="openai/gpt-4o",
|
|
web_search=False,
|
|
cheapest=False,
|
|
fastest=False,
|
|
temperature=1.0,
|
|
pdf_engine="auto",
|
|
chat_mode=False,
|
|
request_timeout=request_timeout,
|
|
)
|
|
|
|
def test_main_openrouter_request_uses_configured_timeout(self):
|
|
response = Mock()
|
|
response.raise_for_status = Mock()
|
|
response.json.return_value = {
|
|
"choices": [{"message": {"content": "done"}}],
|
|
"usage": {"prompt_tokens": 2, "completion_tokens": 3},
|
|
}
|
|
self.node_module.requests.post.return_value = response
|
|
|
|
result = self.call_generate_response(request_timeout=45)
|
|
|
|
self.assertEqual(result[0], "done")
|
|
self.node_module.requests.post.assert_called_once()
|
|
self.assertEqual(self.node_module.requests.post.call_args.kwargs["timeout"], 45)
|
|
|
|
def test_timeout_exception_returns_clear_error(self):
|
|
self.node_module.requests.post.side_effect = FakeTimeout("request timed out")
|
|
|
|
result = self.call_generate_response(request_timeout=2)
|
|
|
|
self.assertIn("API Request Error", result[0])
|
|
self.assertIn("request timed out", result[0])
|
|
self.assertEqual(result[2], "Stats N/A due to error")
|
|
self.assertEqual(result[3], "Credits N/A due to error")
|
|
|
|
def test_fetch_credits_uses_configured_timeout(self):
|
|
self.node.fetch_credits = self.node_module.OpenRouterNode().fetch_credits
|
|
response = Mock()
|
|
response.raise_for_status = Mock()
|
|
response.json.return_value = {
|
|
"data": {
|
|
"total_credits": 10.0,
|
|
"total_usage": 3.25,
|
|
}
|
|
}
|
|
self.node_module.requests.get.return_value = response
|
|
|
|
credits = self.node.fetch_credits("test-key", timeout=45)
|
|
|
|
self.assertEqual(credits, "Remaining: $6.750")
|
|
self.node_module.requests.get.assert_called_once()
|
|
self.assertEqual(self.node_module.requests.get.call_args.kwargs["timeout"], 45)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|