import json import os import sys from unittest.mock import MagicMock import click import pytest from lightning_app import _PROJECT_ROOT from lightning_app.cli.commands.connection import ( _list_app_commands, _PPID, _resolve_command_path, _retrieve_connection_to_an_app, connect, disconnect, ) from lightning_app.utilities import cli_helpers from lightning_app.utilities.commands import base def test_connect_disconnect_local(monkeypatch): disconnect() with pytest.raises(Exception, match="The commands weren't found. Is your app localhost running ?"): connect("localhost", True) with open(os.path.join(os.path.dirname(__file__), "jsons/connect_1.json")) as f: data = json.load(f) data["paths"]["/command/command_with_client"]["post"]["cls_path"] = os.path.join( _PROJECT_ROOT, data["paths"]["/command/command_with_client"]["post"]["cls_path"], ) messages = [] disconnect() def fn(msg): messages.append(msg) monkeypatch.setattr(click, "echo", fn) response = MagicMock() response.status_code = 200 response.json.return_value = data monkeypatch.setattr(cli_helpers.requests, "get", MagicMock(return_value=response)) connect("localhost", True) assert _retrieve_connection_to_an_app() == ("localhost", None) command_path = _resolve_command_path("nested_command") assert not os.path.exists(command_path) command_path = _resolve_command_path("command_with_client") assert os.path.exists(command_path) home = os.path.expanduser("~") s = "/" if sys.platform != "win32" else "\\" command_folder_path = f"{home}{s}.lightning{s}lightning_connection{s}{_PPID}{s}commands" expected = [ f"Find the `command with client` command under {command_folder_path}{s}command_with_client.py.", f"You can review all the downloaded commands under {command_folder_path} folder.", "You are connected to the local Lightning App.", ] assert messages == expected messages = [] connect("localhost", True) assert messages == ["You are connected to the local Lightning App."] messages = [] disconnect() assert messages == ["You are disconnected from the local Lightning App."] messages = [] disconnect() assert messages == [ "You aren't connected to any Lightning App. Please use `lightning connect app_name_or_id` to connect to one." ] assert _retrieve_connection_to_an_app() == (None, None) def test_connect_disconnect_cloud(monkeypatch): disconnect() target_file = _resolve_command_path("command_with_client") if os.path.exists(target_file): os.remove(target_file) with open(os.path.join(os.path.dirname(__file__), "jsons/connect_1.json")) as f: data = json.load(f) data["paths"]["/command/command_with_client"]["post"]["cls_path"] = os.path.join( _PROJECT_ROOT, data["paths"]["/command/command_with_client"]["post"]["cls_path"], ) messages = [] def fn(msg): messages.append(msg) monkeypatch.setattr(click, "echo", fn) response = MagicMock() response.status_code = 200 response.json.return_value = data monkeypatch.setattr(cli_helpers.requests, "get", MagicMock(return_value=response)) project = MagicMock() project.project_id = "custom_project_name" monkeypatch.setattr(cli_helpers, "_get_project", MagicMock(return_value=project)) client = MagicMock() lightningapps = MagicMock() app = MagicMock() app.name = "example" app.id = "1234" lightningapps.lightningapps = [app] client.lightningapp_instance_service_list_lightningapp_instances.return_value = lightningapps monkeypatch.setattr(cli_helpers, "LightningClient", MagicMock(return_value=client)) monkeypatch.setattr(base, "_get_project", MagicMock(return_value=project)) artifact = MagicMock() artifact.filename = "commands/command_with_client.py" artifacts = MagicMock() artifacts.artifacts = [artifact] client.lightningapp_instance_service_list_lightningapp_instance_artifacts.return_value = artifacts monkeypatch.setattr(base, "LightningClient", MagicMock(return_value=client)) with open(data["paths"]["/command/command_with_client"]["post"]["cls_path"], "rb") as f: response.content = f.read() connect("example", True) assert _retrieve_connection_to_an_app() == ("example", "1234") commands = _list_app_commands() assert commands == ["command with client", "command without client", "nested command"] command_path = _resolve_command_path("nested_command") assert not os.path.exists(command_path) command_path = _resolve_command_path("command_with_client") assert os.path.exists(command_path) home = os.path.expanduser("~") s = "/" if sys.platform != "win32" else "\\" command_folder_path = f"{home}{s}.lightning{s}lightning_connection{s}{_PPID}{s}commands" expected = [ f"Storing `command with client` under {command_folder_path}{s}command_with_client.py", f"You can review all the downloaded commands under {command_folder_path} folder.", " ", "The client interface has been successfully installed. ", "You can now run the following commands:", " lightning command without client", " lightning command with client", " lightning nested command", " ", "You are connected to the cloud Lightning App: example.", "Usage: lightning [OPTIONS] COMMAND [ARGS]...", "", " --help Show this message and exit.", "", "Lightning App Commands", " command with client A command with a client.", " command without client A command without a client.", " nested command A nested command.", ] assert messages == expected messages = [] connect("example", True) assert messages == ["You are already connected to the cloud Lightning App: example."] messages = [] disconnect() assert messages == ["You are disconnected from the cloud Lightning App: example."] messages = [] disconnect() assert messages == [ "You aren't connected to any Lightning App. Please use `lightning connect app_name_or_id` to connect to one." ] assert _retrieve_connection_to_an_app() == (None, None)