Meisoftcoltd
75ab5d9676
- Created tests/test_vram_node.py - Mocked dependencies (torch, comfy) to allow testing in isolation - Implemented test cases for zero values, exact GB, fractional values, small values (rounding check), and large values. - Verified test robustness by confirming failure when logic was intentionally altered.
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import sys
|
|
import types
|
|
import pytest
|
|
from unittest.mock import MagicMock
|
|
import os
|
|
|
|
# Mock dependencies before importing anything that might depend on them
|
|
mock_torch = types.ModuleType("torch")
|
|
mock_torch.cuda = MagicMock()
|
|
mock_torch.cuda.is_available.return_value = False
|
|
sys.modules["torch"] = mock_torch
|
|
|
|
sys.modules["comfy"] = types.ModuleType("comfy")
|
|
sys.modules["comfy.model_management"] = types.ModuleType("comfy.model_management")
|
|
|
|
def get_format_bytes():
|
|
# Read the file
|
|
vram_node_path = os.path.join(os.path.dirname(__file__), "..", "vram_node.py")
|
|
with open(vram_node_path, 'r') as f:
|
|
code = f.read()
|
|
|
|
# Mock 'from . import register_node' by replacing it
|
|
code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c")
|
|
|
|
vram_module = types.ModuleType("vram_node")
|
|
exec(code, vram_module.__dict__)
|
|
|
|
return vram_module.format_bytes
|
|
|
|
@pytest.fixture
|
|
def format_bytes():
|
|
return get_format_bytes()
|
|
|
|
def test_format_bytes_zero(format_bytes):
|
|
assert format_bytes(0) == "0.00 GB"
|
|
|
|
def test_format_bytes_gigabyte(format_bytes):
|
|
assert format_bytes(1024**3) == "1.00 GB"
|
|
|
|
def test_format_bytes_fractional(format_bytes):
|
|
assert format_bytes(1024**3 * 0.5) == "0.50 GB"
|
|
assert format_bytes(1024**3 * 2.5) == "2.50 GB"
|
|
|
|
def test_format_bytes_small_values(format_bytes):
|
|
# 1 MB is 1/1024 GB ≈ 0.000976 GB -> 0.00 GB
|
|
assert format_bytes(1024**2) == "0.00 GB"
|
|
# 10 MB is 10/1024 GB ≈ 0.009765 GB -> 0.01 GB (due to rounding .2f)
|
|
assert format_bytes(1024**2 * 10) == "0.01 GB"
|
|
|
|
def test_format_bytes_large_values(format_bytes):
|
|
assert format_bytes(1024**3 * 1024) == "1024.00 GB"
|
|
assert format_bytes(1024**3 * 123.456) == "123.46 GB"
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|