oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2533: added tests + log level for output

Browse Source
This commit is contained in:
Czar Echavez
2024-04-30 15:43:25 +01:00
parent c94fa653f7
commit 59990813f5
8 changed files with 294 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/primaite/__init__.py
View File

@@ -123,9 +123,8 @@ PRIMAITE_PATHS: Final[_PrimaitePaths] = _PrimaitePaths()
def _host_primaite_config() -> None:
if not PRIMAITE_PATHS.app_config_file_path.exists():
pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml"))
shutil.copy2(pkg_config_path, PRIMAITE_PATHS.app_config_file_path)
pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml"))
shutil.copy(pkg_config_path, PRIMAITE_PATHS.app_config_file_path)
_host_primaite_config()

14
src/primaite/cli.py
View File

@@ -9,9 +9,8 @@ import typer
import yaml
from typing_extensions import Annotated
from primaite import _PRIMAITE_ROOT, PRIMAITE_PATHS
from primaite import PRIMAITE_PATHS
from primaite.utils.cli import dev_cli
from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, update_primaite_config
app = typer.Typer(no_args_is_help=True)
app.add_typer(dev_cli.dev, name="dev-mode")
@@ -116,15 +115,4 @@ def setup(overwrite_existing: bool = True) -> None:
_LOGGER.info("Rebuilding the example notebooks...")
reset_example_configs.run(overwrite_existing=True)
_LOGGER.info("Setting default simulation output")
config_dict = get_primaite_config_dict()
if config_dict is None:
return
config_dict["developer_mode"]["output_dir"] = str(_PRIMAITE_ROOT.parent.parent / "simulation_output")
print(f"PrimAITE dev-mode config updated output_dir {config_dict['developer_mode']['output_dir']}")
# update application config
update_primaite_config(config_dict)
_LOGGER.info("PrimAITE setup complete!")

1
src/primaite/setup/_package_data/primaite_config.yaml
View File

@@ -2,6 +2,7 @@
developer_mode:
enabled: False # not enabled by default
sys_log_level: DEBUG # level of output for system logs, DEBUG by default
output_sys_logs: False # system logs not output by default
output_pcap_logs: False # pcap logs not output by default
output_to_terminal: False # do not output to terminal by default

2
src/primaite/simulator/__init__.py
View File

@@ -48,7 +48,7 @@ class _SimOutput:
return self._path
if is_dev_mode():
dev_config = get_primaite_config_dict().get("developer_mode")
return Path(dev_config["output_path"]) if dev_config["output_path"] else self._default_path
return Path(dev_config["output_dir"]) if dev_config["output_dir"] else self._default_path
@path.setter
def path(self, new_path: Path) -> None:

108
src/primaite/utils/cli/dev_cli.py
View File

@@ -1,8 +1,12 @@
from typing import List
import click
import typer
from rich import print
from typing_extensions import Annotated
from primaite import _PRIMAITE_ROOT
from primaite.simulator import LogLevel
from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode, update_primaite_config
dev = typer.Typer()
@@ -69,19 +73,44 @@ def disable():
def config_callback(
ctx: typer.Context,
output_sys_logs: Annotated[
bool, typer.Option("--output-sys-logs/--no-sys-logs", "-sys/-nsys", help="Output system logs to file.")
] = None,
output_pcap_logs: Annotated[
bool,
typer.Option(
"--output-pcap-logs/--no-pcap-logs", "-pcap/-npcap", help="Output network packet capture logs to file."
),
] = None,
output_to_terminal: Annotated[
bool, typer.Option("--output-to_terminal/--no-terminal", "-t/-nt", help="Output system logs to terminal.")
] = None,
ctx: typer.Context,
sys_log_level: Annotated[
LogLevel,
typer.Option(
"--sys-log-level",
"-level",
click_type=click.Choice(LogLevel._member_names_, case_sensitive=False),
help="The level of system logs to output.",
show_default=False
)
] = None,
output_sys_logs: Annotated[
bool,
typer.Option(
"--output-sys-logs/--no-sys-logs",
"-sys/-nsys",
help="Output system logs to file.",
show_default=False
)
] = None,
output_pcap_logs: Annotated[
bool,
typer.Option(
"--output-pcap-logs/--no-pcap-logs",
"-pcap/-npcap",
help="Output network packet capture logs to file.",
show_default=False
),
] = None,
output_to_terminal: Annotated[
bool,
typer.Option(
"--output-to-terminal/--no-terminal",
"-t/-nt",
help="Output system logs to terminal.",
show_default=False
)
] = None,
):
"""Configure the development tools and environment."""
config_dict = get_primaite_config_dict()
@@ -89,6 +118,11 @@ def config_callback(
if config_dict is None:
return
if ctx.params.get("sys_log_level") is not None:
config_dict["developer_mode"]["sys_log_level"] = ctx.params.get("sys_log_level")
print(f"PrimAITE dev-mode config updated sys_log_level={ctx.params.get('sys_log_level')}")
if output_sys_logs is not None:
config_dict["developer_mode"]["output_sys_logs"] = output_sys_logs
print(f"PrimAITE dev-mode config updated {output_sys_logs=}")
@@ -105,28 +139,33 @@ def config_callback(
update_primaite_config(config_dict)
config_typer = typer.Typer(callback=config_callback, name="config", invoke_without_command=True)
config_typer = typer.Typer(
callback=config_callback,
name="config",
no_args_is_help=True,
invoke_without_command=True,
)
dev.add_typer(config_typer)
@config_typer.command()
def path(
directory: Annotated[
str,
typer.Argument(
help="Directory where the system logs and PCAP logs will be output. By default, this will be where the"
"root of the PrimAITE repository is located.",
show_default=False,
),
] = None,
default: Annotated[
bool,
typer.Option(
"--default",
"-root",
help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.",
),
] = None,
directory: Annotated[
str,
typer.Argument(
help="Directory where the system logs and PCAP logs will be output. By default, this will be where the"
"root of the PrimAITE repository is located.",
show_default=False,
),
] = None,
default: Annotated[
bool,
typer.Option(
"--default",
"-root",
help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.",
),
] = None,
):
"""Set the output directory for the PrimAITE system and PCAP logs."""
config_dict = get_primaite_config_dict()
@@ -136,16 +175,15 @@ def path(
if default:
config_dict["developer_mode"]["output_dir"] = None
print(
f"PrimAITE dev-mode config updated output directory will be in "
f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}"
)
# update application config
update_primaite_config(config_dict)
print(f"PrimAITE dev-mode output_dir [medium_turquoise]"
f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}"
f"[/medium_turquoise]")
return
if directory:
config_dict["developer_mode"]["output_dir"] = directory
print(f"PrimAITE dev-mode config updated output_dir={directory}")
# update application config
update_primaite_config(config_dict)
print(f"PrimAITE dev-mode output_dir [medium_turquoise]{directory}[/medium_turquoise]")

25
src/primaite/utils/cli/primaite_config_utils.py
View File

@@ -1,17 +1,30 @@
from typing import Dict
from pathlib import Path
from typing import Dict, Optional
import yaml
from primaite import PRIMAITE_PATHS
def get_primaite_config_dict() -> Dict:
"""Returns a dict containing the PrimAITE application config."""
if PRIMAITE_PATHS.app_config_file_path.exists():
with open(PRIMAITE_PATHS.app_config_file_path, "r") as file:
def get_primaite_config_dict(
config_path: Optional[Path] = None
) -> Dict:
"""
Returns a dict containing the PrimAITE application config.
:param: config_path: takes in a path object - leave empty to use the default app config path
"""
err_msg = "PrimAITE application config could not be loaded."
if config_path is None:
config_path = PRIMAITE_PATHS.app_config_file_path
err_msg = "PrimAITE application config was not found. Have you run `primaite setup`?"
if config_path.exists():
with open(config_path, "r") as file:
return yaml.safe_load(file)
else:
print("PrimAITE application config was not found. Have you run [bold red]primaite setup[/bold red]?")
print(err_msg)
def is_dev_mode() -> bool:

11
tests/integration_tests/cli/__init__.py Normal file
View File

@@ -0,0 +1,11 @@
from typing import List
from typer.testing import CliRunner, Result
from primaite.cli import app
def cli(args: List[str]) -> Result:
"""Pass in a list of arguments and it will return the result."""
runner = CliRunner()
return runner.invoke(app, args)

186
tests/integration_tests/cli/test_dev_cli.py Normal file
View File

@@ -0,0 +1,186 @@
import os
import shutil
import tempfile
from pathlib import Path
import pytest
from primaite import PRIMAITE_PATHS, _PRIMAITE_ROOT
from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict
from tests.integration_tests.cli import cli
@pytest.fixture(autouse=True)
def test_setup():
"""
Setup this test by copying the
"""
original_config_path = PRIMAITE_PATHS.app_config_file_path # keep track of app config before test
temp_dir = tempfile.gettempdir()
temp_config = Path(temp_dir) / "primaite_config.yaml"
shutil.copyfile(
_PRIMAITE_ROOT / "setup" / "_package_data" / "primaite_config.yaml", temp_config
) # copy the default primaite config to temp directory
PRIMAITE_PATHS.app_config_file_path = temp_config # use the copy for the test
yield # run test
os.remove(temp_config) # clean up temp file
PRIMAITE_PATHS.app_config_file_path = original_config_path # restore app conf because other devs will yell at me
def test_dev_mode_enable_disable():
"""Test dev mode enable and disable."""
# check defaults
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["enabled"] is False # not enabled by default
result = cli(["dev-mode", "show"])
assert "Production" in result.output # should print that it is in Production mode by default
result = cli(["dev-mode", "enable"])
assert "Development" in result.output # should print that it is in Development mode
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["enabled"] # config should reflect that dev mode is enabled
result = cli(["dev-mode", "show"])
assert "Development" in result.output # should print that it is in Development mode
result = cli(["dev-mode", "disable"])
assert "Production" in result.output # should print that it is in Production mode
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["enabled"] is False # config should reflect that dev mode is disabled
result = cli(["dev-mode", "show"])
assert "Production" in result.output # should print that it is in Production mode
def test_dev_mode_config_sys_log_level():
"""Check that the system log level can be changed via CLI."""
# check defaults
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["sys_log_level"] == "DEBUG" # DEBUG by default
result = cli(["dev-mode", "config", "-level", "WARNING"])
assert "sys_log_level=WARNING" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that log level is WARNING
assert config_dict["developer_mode"]["sys_log_level"] == "WARNING"
result = cli(["dev-mode", "config", "--sys-log-level", "INFO"])
assert "sys_log_level=INFO" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that log level is WARNING
assert config_dict["developer_mode"]["sys_log_level"] == "INFO"
def test_dev_mode_config_sys_logs_enable_disable():
"""Test that the system logs output can be enabled or disabled."""
# check defaults
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["output_sys_logs"] is False # False by default
result = cli(["dev-mode", "config", "--output-sys-logs"])
assert "output_sys_logs=True" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_sys_logs is True
assert config_dict["developer_mode"]["output_sys_logs"]
result = cli(["dev-mode", "config", "--no-sys-logs"])
assert "output_sys_logs=False" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_sys_logs is True
assert config_dict["developer_mode"]["output_sys_logs"] is False
result = cli(["dev-mode", "config", "-sys"])
assert "output_sys_logs=True" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_sys_logs is True
assert config_dict["developer_mode"]["output_sys_logs"]
result = cli(["dev-mode", "config", "-nsys"])
assert "output_sys_logs=False" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_sys_logs is True
assert config_dict["developer_mode"]["output_sys_logs"] is False
def test_dev_mode_config_pcap_logs_enable_disable():
"""Test that the pcap logs output can be enabled or disabled."""
# check defaults
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["output_pcap_logs"] is False # False by default
result = cli(["dev-mode", "config", "--output-pcap-logs"])
assert "output_pcap_logs=True" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_pcap_logs is True
assert config_dict["developer_mode"]["output_pcap_logs"]
result = cli(["dev-mode", "config", "--no-pcap-logs"])
assert "output_pcap_logs=False" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_pcap_logs is True
assert config_dict["developer_mode"]["output_pcap_logs"] is False
result = cli(["dev-mode", "config", "-pcap"])
assert "output_pcap_logs=True" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_pcap_logs is True
assert config_dict["developer_mode"]["output_pcap_logs"]
result = cli(["dev-mode", "config", "-npcap"])
assert "output_pcap_logs=False" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_pcap_logs is True
assert config_dict["developer_mode"]["output_pcap_logs"] is False
def test_dev_mode_config_output_to_terminal_enable_disable():
"""Test that the output to terminal can be enabled or disabled."""
# check defaults
config_dict = get_primaite_config_dict()
assert config_dict["developer_mode"]["output_to_terminal"] is False # False by default
result = cli(["dev-mode", "config", "--output-to-terminal"])
assert "output_to_terminal=True" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_to_terminal is True
assert config_dict["developer_mode"]["output_to_terminal"]
result = cli(["dev-mode", "config", "--no-terminal"])
assert "output_to_terminal=False" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_to_terminal is True
assert config_dict["developer_mode"]["output_to_terminal"] is False
result = cli(["dev-mode", "config", "-t"])
assert "output_to_terminal=True" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_to_terminal is True
assert config_dict["developer_mode"]["output_to_terminal"]
result = cli(["dev-mode", "config", "-nt"])
assert "output_to_terminal=False" in result.output # should print correct value
config_dict = get_primaite_config_dict()
# config should reflect that output_to_terminal is True
assert config_dict["developer_mode"]["output_to_terminal"] is False