|
1 | 1 | import socket |
2 | 2 | import threading |
3 | 3 | import time |
4 | | -from unittest.mock import patch |
| 4 | +from unittest.mock import MagicMock, patch |
5 | 5 |
|
6 | 6 | import pytest |
7 | 7 | import typer |
8 | | -from langflow.__main__ import _create_superuser, app, get_number_of_workers |
| 8 | +from langflow.__main__ import _create_superuser, api_key_banner, app, get_number_of_workers |
9 | 9 | from lfx.services import deps |
10 | 10 |
|
11 | 11 |
|
@@ -147,6 +147,101 @@ async def test_failed_auth_token_validation(self, client, active_super_user): # |
147 | 147 | assert exc_info.value.exit_code == 1 |
148 | 148 |
|
149 | 149 |
|
| 150 | +class TestApiKeyBanner: |
| 151 | + """Tests for api_key_banner clipboard fallback (headless environments).""" |
| 152 | + |
| 153 | + def _make_key(self, value: str = "sk-test-1234"): |
| 154 | + mock = MagicMock() |
| 155 | + mock.api_key = value |
| 156 | + return mock |
| 157 | + |
| 158 | + def test_clipboard_available_copies_key(self): |
| 159 | + """When pyperclip works, key is copied and clipboard hint is shown.""" |
| 160 | + key = self._make_key() |
| 161 | + with ( |
| 162 | + patch("pyperclip.copy") as mock_copy, |
| 163 | + patch("langflow.__main__.Console") as mock_console_cls, |
| 164 | + ): |
| 165 | + mock_console = MagicMock() |
| 166 | + mock_console_cls.return_value = mock_console |
| 167 | + |
| 168 | + api_key_banner(key) |
| 169 | + |
| 170 | + mock_copy.assert_called_once_with(key.api_key) |
| 171 | + printed = mock_console.print.call_args[0][0] |
| 172 | + assert "clipboard" in str(printed).lower() |
| 173 | + assert key.api_key in str(printed) |
| 174 | + |
| 175 | + def test_clipboard_unavailable_still_prints_key(self): |
| 176 | + """When pyperclip raises (headless/Docker), key is still displayed on stdout.""" |
| 177 | + key = self._make_key() |
| 178 | + with ( |
| 179 | + patch("pyperclip.copy", side_effect=Exception("No clipboard mechanism")), |
| 180 | + patch("langflow.__main__.Console") as mock_console_cls, |
| 181 | + ): |
| 182 | + mock_console = MagicMock() |
| 183 | + mock_console_cls.return_value = mock_console |
| 184 | + |
| 185 | + # Must NOT raise |
| 186 | + api_key_banner(key) |
| 187 | + |
| 188 | + mock_console.print.assert_called_once() |
| 189 | + printed = mock_console.print.call_args[0][0] |
| 190 | + assert key.api_key in str(printed) |
| 191 | + |
| 192 | + def test_clipboard_unavailable_shows_fallback_hint(self): |
| 193 | + """When clipboard is unavailable, hint text must not mention clipboard.""" |
| 194 | + key = self._make_key() |
| 195 | + with ( |
| 196 | + patch("pyperclip.copy", side_effect=Exception("No clipboard mechanism")), |
| 197 | + patch("langflow.__main__.Console") as mock_console_cls, |
| 198 | + ): |
| 199 | + mock_console = MagicMock() |
| 200 | + mock_console_cls.return_value = mock_console |
| 201 | + |
| 202 | + api_key_banner(key) |
| 203 | + |
| 204 | + printed = str(mock_console.print.call_args[0][0]) |
| 205 | + assert "clipboard" not in printed.lower() |
| 206 | + assert "securely" in printed.lower() |
| 207 | + |
| 208 | + def test_unicode_error_fallback_with_clipboard(self): |
| 209 | + """On UnicodeEncodeError, logger fallback includes clipboard message when available.""" |
| 210 | + key = self._make_key() |
| 211 | + with ( |
| 212 | + patch("pyperclip.copy"), |
| 213 | + patch("langflow.__main__.Console") as mock_console_cls, |
| 214 | + patch("langflow.__main__.logger") as mock_logger, |
| 215 | + ): |
| 216 | + mock_console = MagicMock() |
| 217 | + mock_console.print.side_effect = UnicodeEncodeError("utf-8", b"", 0, 1, "reason") |
| 218 | + mock_console_cls.return_value = mock_console |
| 219 | + |
| 220 | + api_key_banner(key) |
| 221 | + |
| 222 | + logged_messages = " ".join(str(c) for c in mock_logger.info.call_args_list) |
| 223 | + assert key.api_key in logged_messages |
| 224 | + assert "clipboard" in logged_messages.lower() |
| 225 | + |
| 226 | + def test_unicode_error_fallback_without_clipboard(self): |
| 227 | + """On UnicodeEncodeError, logger fallback omits clipboard message when unavailable.""" |
| 228 | + key = self._make_key() |
| 229 | + with ( |
| 230 | + patch("pyperclip.copy", side_effect=Exception("No clipboard mechanism")), |
| 231 | + patch("langflow.__main__.Console") as mock_console_cls, |
| 232 | + patch("langflow.__main__.logger") as mock_logger, |
| 233 | + ): |
| 234 | + mock_console = MagicMock() |
| 235 | + mock_console.print.side_effect = UnicodeEncodeError("utf-8", b"", 0, 1, "reason") |
| 236 | + mock_console_cls.return_value = mock_console |
| 237 | + |
| 238 | + api_key_banner(key) |
| 239 | + |
| 240 | + logged_messages = " ".join(str(c) for c in mock_logger.info.call_args_list) |
| 241 | + assert key.api_key in logged_messages |
| 242 | + assert "clipboard" not in logged_messages.lower() |
| 243 | + |
| 244 | + |
150 | 245 | def test_get_number_of_workers(): |
151 | 246 | """Test that get_number_of_workers uses cpu_count on Linux.""" |
152 | 247 | with ( |
|
0 commit comments