133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
import nomad_client
|
|
import subprocess
|
|
|
|
@patch("subprocess.run")
|
|
@patch("nomad_client.get_node_map")
|
|
def test_get_allocation_id(mock_node_map, mock_run):
|
|
"""Test getting allocation ID for a node."""
|
|
mock_node_map.return_value = {"node_id1": "node1"}
|
|
|
|
# Mock 'nomad job status navidrome-litefs' output
|
|
mock_job_status = MagicMock()
|
|
mock_job_status.stdout = """
|
|
Allocations
|
|
ID Node ID Task Group Version Desired Status Created Modified
|
|
abc12345 node_id1 navidrome 1 run running 1h ago 1h ago
|
|
"""
|
|
|
|
# Mock 'nomad alloc status abc12345' output
|
|
mock_alloc_status = MagicMock()
|
|
mock_alloc_status.stdout = "ID = abc12345-full-id"
|
|
|
|
mock_run.side_effect = [mock_job_status, mock_alloc_status]
|
|
|
|
alloc_id = nomad_client.get_allocation_id("node1", "navidrome-litefs")
|
|
assert alloc_id == "abc12345-full-id"
|
|
|
|
@patch("subprocess.run")
|
|
def test_get_logs(mock_run):
|
|
"""Test fetching logs for an allocation."""
|
|
mock_stderr = "Error: database is locked\nSome other error"
|
|
m = MagicMock()
|
|
m.stdout = mock_stderr
|
|
m.return_code = 0
|
|
mock_run.return_value = m
|
|
|
|
logs = nomad_client.get_allocation_logs("abc12345", tail=20)
|
|
assert "database is locked" in logs
|
|
# It should have tried with -task navidrome first
|
|
mock_run.assert_any_call(
|
|
["nomad", "alloc", "logs", "-stderr", "-task", "navidrome", "-n", "20", "abc12345"],
|
|
capture_output=True, text=True, check=True
|
|
)
|
|
|
|
@patch("subprocess.run")
|
|
def test_restart_allocation(mock_run):
|
|
"""Test restarting an allocation."""
|
|
m = MagicMock()
|
|
m.return_code = 0
|
|
mock_run.return_value = m
|
|
|
|
success = nomad_client.restart_allocation("abc12345")
|
|
assert success is True
|
|
mock_run.assert_called_with(
|
|
["nomad", "alloc", "restart", "abc12345"],
|
|
capture_output=True, text=True, check=True
|
|
)
|
|
|
|
@patch("subprocess.run")
|
|
def test_exec_command(mock_run):
|
|
"""Test executing a command in an allocation."""
|
|
m = MagicMock()
|
|
m.stdout = "Command output"
|
|
m.return_code = 0
|
|
mock_run.return_value = m
|
|
|
|
output = nomad_client.exec_command("abc12345", ["ls", "/data"])
|
|
assert output == "Command output"
|
|
mock_run.assert_called_with(
|
|
["nomad", "alloc", "exec", "-task", "navidrome", "abc12345", "ls", "/data"],
|
|
capture_output=True, text=True, check=True
|
|
)
|
|
|
|
@patch("subprocess.run")
|
|
def test_exec_command_failure(mock_run):
|
|
"""Test executing a command handles failure gracefully."""
|
|
mock_run.side_effect = subprocess.CalledProcessError(1, "nomad", stderr="Nomad error")
|
|
|
|
output = nomad_client.exec_command("abc12345", ["ls", "/data"])
|
|
assert "Nomad Error" in output
|
|
assert "Nomad error" not in output # The exception str might not contain stderr directly depending on python version
|
|
|
|
@patch("subprocess.run")
|
|
def test_get_node_map_failure(mock_run):
|
|
"""Test get_node_map handles failure."""
|
|
mock_run.side_effect = FileNotFoundError("No such file")
|
|
|
|
# It should not raise
|
|
node_map = nomad_client.get_node_map()
|
|
assert node_map == {}
|
|
|
|
@patch("subprocess.run")
|
|
def test_get_job_allocations(mock_run):
|
|
"""Test getting all allocations for a job with their IPs."""
|
|
# Mock 'nomad job status navidrome-litefs'
|
|
mock_job_status = MagicMock()
|
|
mock_job_status.stdout = """
|
|
Allocations
|
|
ID Node ID Node Name Status Created
|
|
abc12345 node1 host1 running 1h ago
|
|
def67890 node2 host2 running 1h ago
|
|
"""
|
|
|
|
# Mock 'nomad alloc status' for each alloc
|
|
mock_alloc1 = MagicMock()
|
|
mock_alloc1.stdout = """
|
|
ID = abc12345
|
|
Node Name = host1
|
|
Allocation Addresses:
|
|
Label Dynamic Address
|
|
*http yes 1.1.1.1:4533 -> 4533
|
|
*litefs yes 1.1.1.1:20202 -> 20202
|
|
"""
|
|
mock_alloc2 = MagicMock()
|
|
mock_alloc2.stdout = """
|
|
ID = def67890
|
|
Node Name = host2
|
|
Allocation Addresses:
|
|
Label Dynamic Address
|
|
*http yes 2.2.2.2:4533 -> 4533
|
|
*litefs yes 2.2.2.2:20202 -> 20202
|
|
"""
|
|
|
|
mock_run.side_effect = [mock_job_status, mock_alloc1, mock_alloc2]
|
|
|
|
# This should fail initially because nomad_client.get_job_allocations doesn't exist
|
|
try:
|
|
allocs = nomad_client.get_job_allocations("navidrome-litefs")
|
|
assert len(allocs) == 2
|
|
assert allocs[0]["ip"] == "1.1.1.1"
|
|
except AttributeError:
|
|
pytest.fail("nomad_client.get_job_allocations not implemented") |