File size: 5,450 Bytes
6a95a78
 
 
 
1c0aad9
6a95a78
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
905e1ea
 
6a95a78
 
 
 
 
1c0aad9
6a95a78
 
 
 
 
 
 
 
 
 
1c0aad9
 
 
 
 
6a95a78
 
 
 
1c0aad9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6a95a78
 
1c0aad9
6a95a78
 
 
 
 
1c0aad9
 
6a95a78
 
 
1c0aad9
6a95a78
 
1c0aad9
6a95a78
 
 
 
 
 
 
 
 
 
 
1c0aad9
6a95a78
 
 
 
 
 
 
 
 
 
 
 
 
1c0aad9
6a95a78
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import pytest
from unittest.mock import MagicMock, patch

from blossomtune_gradio import processing
from blossomtune_gradio.database import Config


@pytest.fixture(autouse=True)
def reset_process_store():
    """
    Fixture that runs automatically for every test in this module.
    It resets the process_store to its default state to ensure test isolation.
    """
    processing.process_store = {"superlink": None, "runner": None}
    yield
    processing.process_store = {"superlink": None, "runner": None}


@patch("blossomtune_gradio.processing.threading.Thread")
@patch(
    "blossomtune_gradio.processing.shutil.which",
    return_value="/fake/path/flower-superlink",
)
def test_start_superlink_success(mock_which, mock_thread):
    """Verify that start_superlink successfully starts a thread."""
    success, message = processing.start_superlink()

    assert success is True
    assert message == "Superlink process started."
    mock_thread.assert_called_once()
    call_args = mock_thread.call_args
    assert call_args.kwargs["target"] == processing.run_process
    assert call_args.kwargs["args"][0][0] == "/fake/path/flower-superlink"
    assert call_args.kwargs["args"][0][1] == "--ssl-ca-certfile"


def test_start_superlink_already_running(mocker):
    """Verify that start_superlink returns False if a process is already running."""
    mock_process = MagicMock()
    mock_process.poll.return_value = None
    processing.process_store["superlink"] = mock_process

    success, message = processing.start_superlink()
    assert success is False
    assert "already running" in message


@patch("blossomtune_gradio.processing.os.path.exists", return_value=True)
@patch("blossomtune_gradio.processing.threading.Thread")
@patch("blossomtune_gradio.processing.shutil.which", return_value="/fake/path/flwr")
def test_start_runner_success_internal_superlink(
    mock_which, mock_thread, mock_exists, db_session
):
    """Verify start_runner succeeds with an internal superlink and updates the DB."""
    # Arrange: Mock a running internal superlink process
    mock_superlink = MagicMock()
    mock_superlink.poll.return_value = None
    processing.process_store["superlink"] = mock_superlink

    # Act
    success, message = processing.start_runner("app.main", "run1", "15")

    # Assert
    assert success is True
    assert message == "Federation Run is starting...."
    mock_thread.assert_called_once()

    # Verify DB was updated using SQLAlchemy
    config_entry = db_session.query(Config).filter_by(key="num_partitions").first()
    assert config_entry is not None
    assert config_entry.value == "15"


@patch("blossomtune_gradio.processing.os.path.exists", return_value=True)
@patch("blossomtune_gradio.processing.threading.Thread")
@patch("blossomtune_gradio.processing.shutil.which", return_value="/fake/path/flwr")
def test_start_runner_success_external_superlink(
    mock_which, mock_thread, mock_exists, db_session, mocker
):
    """Verify start_runner succeeds with an external superlink."""
    # Arrange
    mocker.patch("blossomtune_gradio.config.SUPERLINK_MODE", "external")
    mocker.patch("blossomtune_gradio.util.is_port_open", return_value=True)

    # Act
    success, message = processing.start_runner("app.main", "run1", "10")

    # Assert
    assert success is True
    assert message == "Federation Run is starting...."
    mock_thread.assert_called_once()


def test_start_runner_internal_superlink_not_running(db_session):
    """Verify start_runner fails if internal superlink is not running."""
    processing.process_store["superlink"] = None
    success, message = processing.start_runner("app.main", "run1", "10")
    assert success is False
    assert "Internal Superlink is not running" in message


def test_start_runner_missing_args(db_session):
    """Verify start_runner fails if arguments are missing."""
    mock_superlink = MagicMock()
    mock_superlink.poll.return_value = None
    processing.process_store["superlink"] = mock_superlink

    success, message = processing.start_runner("", "run1", "10")
    assert success is False
    assert "provide a Runner App" in message


@patch("blossomtune_gradio.processing.os.path.exists", return_value=False)
def test_start_runner_app_path_not_found(mock_exists, db_session):
    """Verify start_runner fails if the app path doesn't exist."""
    mock_superlink = MagicMock()
    mock_superlink.poll.return_value = None
    processing.process_store["superlink"] = mock_superlink

    success, message = processing.start_runner("non.existent.app", "run1", "10")
    assert success is False
    assert "Unable to find app path" in message


def test_stop_process_running():
    """Verify stop_process terminates a running process."""
    mock_process = MagicMock()
    mock_process.poll.return_value = None
    processing.process_store["superlink"] = mock_process

    processing.stop_process("superlink")

    mock_process.terminate.assert_called_once()
    mock_process.wait.assert_called_once()
    assert processing.process_store["superlink"] is None


def test_stop_process_not_running(mocker):
    """Verify stop_process does nothing if the process is not running."""
    log_mock = mocker.patch("blossomtune_gradio.processing.log")
    processing.process_store["superlink"] = None

    processing.stop_process("superlink")

    log_mock.assert_any_call(
        "[Superlink] Stop command received, but no process was running."
    )