95 lines
2.5 KiB
Python
95 lines
2.5 KiB
Python
import os
|
|
import runpy
|
|
import sys
|
|
from unittest import mock
|
|
from unittest.mock import ANY, Mock
|
|
|
|
import pytest
|
|
|
|
from lightning_app import LightningFlow
|
|
from lightning_app.frontend.stream_lit import StreamlitFrontend
|
|
from lightning_app.utilities.state import AppState
|
|
|
|
|
|
def test_stop_server_not_running():
|
|
frontend = StreamlitFrontend(render_fn=Mock())
|
|
with pytest.raises(RuntimeError, match="Server is not running."):
|
|
frontend.stop_server()
|
|
|
|
|
|
def _noop_render_fn(_):
|
|
pass
|
|
|
|
|
|
class MockFlow(LightningFlow):
|
|
@property
|
|
def name(self):
|
|
return "root.my.flow"
|
|
|
|
def run(self):
|
|
pass
|
|
|
|
|
|
@mock.patch("lightning_app.frontend.stream_lit.subprocess")
|
|
def test_streamlit_frontend_start_stop_server(subprocess_mock):
|
|
"""Test that `StreamlitFrontend.start_server()` invokes subprocess.Popen with the right parameters."""
|
|
frontend = StreamlitFrontend(render_fn=_noop_render_fn)
|
|
frontend.flow = MockFlow()
|
|
frontend.start_server(host="hostname", port=1111)
|
|
subprocess_mock.Popen.assert_called_once()
|
|
|
|
env_variables = subprocess_mock.method_calls[0].kwargs["env"]
|
|
call_args = subprocess_mock.method_calls[0].args[0]
|
|
assert call_args == [
|
|
sys.executable,
|
|
"-m",
|
|
"streamlit",
|
|
"run",
|
|
ANY,
|
|
"--server.address",
|
|
"hostname",
|
|
"--server.port",
|
|
"1111",
|
|
"--server.baseUrlPath",
|
|
"root.my.flow",
|
|
"--server.headless",
|
|
"true",
|
|
]
|
|
|
|
assert env_variables["LIGHTNING_FLOW_NAME"] == "root.my.flow"
|
|
assert env_variables["LIGHTNING_RENDER_FUNCTION"] == "_noop_render_fn"
|
|
assert env_variables["LIGHTNING_RENDER_MODULE_FILE"] == __file__
|
|
|
|
assert "LIGHTNING_FLOW_NAME" not in os.environ
|
|
assert "LIGHTNING_RENDER_FUNCTION" not in os.environ
|
|
assert "LIGHTNING_RENDER_MODULE_FILE" not in os.environ
|
|
|
|
frontend.stop_server()
|
|
subprocess_mock.Popen().kill.assert_called_once()
|
|
|
|
|
|
def _streamlit_call_me(state):
|
|
assert isinstance(state, AppState)
|
|
|
|
|
|
@mock.patch.dict(
|
|
os.environ,
|
|
{
|
|
"LIGHTNING_FLOW_NAME": "root",
|
|
"LIGHTNING_RENDER_FUNCTION": "_streamlit_call_me",
|
|
"LIGHTNING_RENDER_MODULE_FILE": __file__,
|
|
},
|
|
)
|
|
def test_streamlit_wrapper_calls_render_fn(*_):
|
|
runpy.run_module("lightning_app.frontend.streamlit_base")
|
|
# TODO: find a way to assert that _streamlit_call_me got called
|
|
|
|
|
|
def test_method_exception():
|
|
class A:
|
|
def render_fn(self):
|
|
pass
|
|
|
|
with pytest.raises(TypeError, match="being a method"):
|
|
StreamlitFrontend(render_fn=A().render_fn)
|