mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
zeekay
b90e7752eb
- Added net_check utility for verifying hanzo/net installation - Improved error handling and dependency detection - Added comprehensive tests for node command - Set up CI/CD workflows for both projects - Better handling of venv and Python paths
163 lines
6.2 KiB
Python
163 lines
6.2 KiB
Python
"""Tests for hanzo node command integration with hanzo/net."""
|
|
|
|
import pytest
|
|
import subprocess
|
|
import sys
|
|
from unittest.mock import patch, MagicMock, call
|
|
from pathlib import Path
|
|
|
|
from hanzo.cli import start_compute_node
|
|
from hanzo.utils.net_check import check_net_installation, get_missing_dependencies
|
|
|
|
|
|
class TestNetCheck:
|
|
"""Test hanzo/net availability checking."""
|
|
|
|
def test_check_net_installation_with_venv(self, tmp_path):
|
|
"""Test checking net installation with venv present."""
|
|
# Create mock net directory with venv
|
|
net_path = tmp_path / "net"
|
|
net_path.mkdir()
|
|
venv_path = net_path / ".venv" / "bin"
|
|
venv_path.mkdir(parents=True)
|
|
(venv_path / "python").touch()
|
|
|
|
with patch("hanzo.utils.net_check.Path") as mock_path:
|
|
mock_path.home.return_value = tmp_path.parent
|
|
mock_path.return_value.exists.return_value = True
|
|
|
|
with patch("subprocess.run") as mock_run:
|
|
mock_run.return_value.returncode = 0
|
|
|
|
is_available, path, python = check_net_installation()
|
|
|
|
assert is_available is True
|
|
assert path is not None
|
|
assert "python" in python
|
|
|
|
def test_check_net_installation_no_venv(self):
|
|
"""Test checking net installation without venv."""
|
|
with patch("hanzo.utils.net_check.Path") as mock_path:
|
|
mock_path.return_value.exists.return_value = False
|
|
|
|
with patch("subprocess.run") as mock_run:
|
|
mock_run.return_value.returncode = 1
|
|
|
|
is_available, path, python = check_net_installation()
|
|
|
|
assert is_available is False
|
|
|
|
def test_get_missing_dependencies(self):
|
|
"""Test checking for missing dependencies."""
|
|
with patch("subprocess.run") as mock_run:
|
|
# Mock some packages as installed, others as missing
|
|
def side_effect(cmd, **kwargs):
|
|
result = MagicMock()
|
|
if "scapy" in cmd[2] or "aiohttp" in cmd[2]:
|
|
result.returncode = 0
|
|
else:
|
|
result.returncode = 1
|
|
return result
|
|
|
|
mock_run.side_effect = side_effect
|
|
|
|
missing = get_missing_dependencies()
|
|
|
|
assert "scapy" not in missing
|
|
assert "aiohttp" not in missing
|
|
assert "mlx" in missing
|
|
assert "transformers" in missing
|
|
|
|
|
|
class TestNodeCommand:
|
|
"""Test hanzo node command."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_compute_node_missing_net(self):
|
|
"""Test starting compute node when net is not available."""
|
|
ctx = MagicMock()
|
|
ctx.obj = {"console": MagicMock()}
|
|
|
|
with patch("hanzo.cli.check_net_installation") as mock_check:
|
|
mock_check.return_value = (False, None, None)
|
|
|
|
await start_compute_node(ctx)
|
|
|
|
# Should print error and return early
|
|
ctx.obj["console"].print.assert_any_call(
|
|
"[red]Error:[/red] hanzo/net is not properly configured"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_compute_node_with_net_venv(self):
|
|
"""Test starting compute node with net venv available."""
|
|
ctx = MagicMock()
|
|
ctx.obj = {"console": MagicMock()}
|
|
|
|
with patch("hanzo.cli.check_net_installation") as mock_check:
|
|
mock_check.return_value = (True, "/path/to/net", "/path/to/net/.venv/bin/python")
|
|
|
|
with patch("subprocess.run") as mock_run:
|
|
mock_run.return_value.returncode = 0
|
|
|
|
with patch("os.chdir"):
|
|
await start_compute_node(ctx, models=("llama-3.2-3b",))
|
|
|
|
# Should use venv python
|
|
mock_run.assert_called_once()
|
|
assert "/path/to/net/.venv/bin/python" in mock_run.call_args[0][0][0]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_compute_node_with_models(self):
|
|
"""Test starting compute node with specific models."""
|
|
ctx = MagicMock()
|
|
ctx.obj = {"console": MagicMock()}
|
|
|
|
with patch("hanzo.cli.check_net_installation") as mock_check:
|
|
mock_check.return_value = (True, "/path/to/net", sys.executable)
|
|
|
|
with patch("subprocess.run") as mock_run:
|
|
mock_run.return_value.returncode = 0
|
|
|
|
with patch("os.chdir"):
|
|
await start_compute_node(
|
|
ctx,
|
|
models=("llama-3.2-3b", "gpt-4"),
|
|
name="test-node",
|
|
port=8080
|
|
)
|
|
|
|
# Check environment variables were set
|
|
env = mock_run.call_args[1]["env"]
|
|
assert env["NET_MODELS"] == "llama-3.2-3b,gpt-4"
|
|
assert env["NET_NODE_NAME"] == "test-node"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_compute_node_keyboard_interrupt(self):
|
|
"""Test graceful shutdown on keyboard interrupt."""
|
|
ctx = MagicMock()
|
|
ctx.obj = {"console": MagicMock()}
|
|
|
|
with patch("hanzo.cli.check_net_installation") as mock_check:
|
|
mock_check.return_value = (True, "/path/to/net", sys.executable)
|
|
|
|
with patch("subprocess.run") as mock_run:
|
|
mock_run.side_effect = KeyboardInterrupt()
|
|
|
|
with patch("os.chdir"):
|
|
await start_compute_node(ctx)
|
|
|
|
# Should print shutdown message
|
|
ctx.obj["console"].print.assert_any_call(
|
|
"\n[yellow]Shutting down node...[/yellow]"
|
|
)
|
|
|
|
|
|
def test_cli_imports():
|
|
"""Test that CLI modules can be imported."""
|
|
from hanzo import cli
|
|
from hanzo.utils import net_check
|
|
|
|
assert hasattr(cli, "start_compute_node")
|
|
assert hasattr(net_check, "check_net_installation")
|
|
assert hasattr(net_check, "get_missing_dependencies") |