1 files changed, 160 insertions, 0 deletions
diff --git a/tests/aiet/test_backend_controller.py b/tests/aiet/test_backend_controller.py
new file mode 100644
index 0000000..8836ec5
--- /dev/null
+++ b/tests/aiet/test_backend_controller.py
@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for system controller."""
+import csv
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import psutil
+import pytest
+
+from aiet.backend.common import ConfigurationException
+from aiet.backend.controller import SystemController
+from aiet.backend.controller import SystemControllerSingleInstance
+from aiet.utils.proc import ShellCommand
+
+
+def get_system_controller(**kwargs: Any) -> SystemController:
+    """Get service controller."""
+    single_instance = kwargs.get("single_instance", False)
+    if single_instance:
+        pid_file_path = kwargs.get("pid_file_path")
+        return SystemControllerSingleInstance(pid_file_path)
+
+    return SystemController()
+
+
+def test_service_controller() -> None:
+    """Test service controller functionality."""
+    service_controller = get_system_controller()
+
+    assert service_controller.get_output() == ("", "")
+    with pytest.raises(ConfigurationException, match="Wrong working directory"):
+        service_controller.start(["sleep 100"], Path("unknown"))
+
+    service_controller.start(["sleep 100"], Path.cwd())
+    assert service_controller.is_running()
+
+    service_controller.stop(True)
+    assert not service_controller.is_running()
+    assert service_controller.get_output() == ("", "")
+
+    service_controller.stop()
+
+    with pytest.raises(
+        ConfigurationException, match="System should have only one command to run"
+    ):
+        service_controller.start(["sleep 100", "sleep 101"], Path.cwd())
+
+    with pytest.raises(ConfigurationException, match="No startup command provided"):
+        service_controller.start([""], Path.cwd())
+
+
+def test_service_controller_bad_configuration() -> None:
+    """Test service controller functionality for bad configuration."""
+    with pytest.raises(Exception, match="No pid file path presented"):
+        service_controller = get_system_controller(
+            single_instance=True, pid_file_path=None
+        )
+        service_controller.start(["sleep 100"], Path.cwd())
+
+
+def test_service_controller_writes_process_info_correctly(tmpdir: Any) -> None:
+    """Test that controller writes process info correctly."""
+    pid_file = Path(tmpdir) / "test.pid"
+
+    service_controller = get_system_controller(
+        single_instance=True, pid_file_path=Path(tmpdir) / "test.pid"
+    )
+
+    service_controller.start(["sleep 100"], Path.cwd())
+    assert service_controller.is_running()
+    assert pid_file.is_file()
+
+    with open(pid_file, "r", encoding="utf-8") as file:
+        csv_reader = csv.reader(file)
+        rows = list(csv_reader)
+        assert len(rows) == 1
+
+        name, *_ = rows[0]
+        assert name == "sleep"
+
+    service_controller.stop()
+    assert pid_file.exists()
+
+
+def test_service_controller_does_not_write_process_info_if_process_finishes(
+    tmpdir: Any,
+) -> None:
+    """Test that controller does not write process info if process already finished."""
+    pid_file = Path(tmpdir) / "test.pid"
+    service_controller = get_system_controller(
+        single_instance=True, pid_file_path=pid_file
+    )
+    service_controller.is_running = lambda: False  # type: ignore
+    service_controller.start(["echo hello"], Path.cwd())
+
+    assert not pid_file.exists()
+
+
+def test_service_controller_searches_for_previous_instances_correctly(
+    tmpdir: Any,
+) -> None:
+    """Test that controller searches for previous instances correctly."""
+    pid_file = Path(tmpdir) / "test.pid"
+    command = ShellCommand().run("sleep", "100")
+    assert command.is_alive()
+
+    pid = command.process.pid
+    process = psutil.Process(pid)
+    with open(pid_file, "w", encoding="utf-8") as file:
+        csv_writer = csv.writer(file)
+        csv_writer.writerow(("some_process", "some_program", "some_cwd", os.getpid()))
+        csv_writer.writerow((process.name(), process.exe(), process.cwd(), process.pid))
+        csv_writer.writerow(("some_old_process", "not_running", "from_nowhere", 77777))
+
+    service_controller = get_system_controller(
+        single_instance=True, pid_file_path=pid_file
+    )
+    service_controller.start(["sleep 100"], Path.cwd())
+    # controller should stop this process as it is currently running and
+    # mentioned in pid file
+    assert not command.is_alive()
+
+    service_controller.stop()
+
+
+@pytest.mark.parametrize(
+    "executable", ["test_backend_run_script.sh", "test_backend_run"]
+)
+def test_service_controller_run_shell_script(
+    executable: str, test_resources_path: Path
+) -> None:
+    """Test controller's ability to run shell scripts."""
+    script_path = test_resources_path / "scripts"
+
+    service_controller = get_system_controller()
+
+    service_controller.start([executable], script_path)
+
+    assert service_controller.is_running()
+    # give time for the command to produce output
+    time.sleep(2)
+    service_controller.stop(wait=True)
+    assert not service_controller.is_running()
+    stdout, stderr = service_controller.get_output()
+    assert stdout == "Hello from script\n"
+    assert stderr == "Oops!\n"
+
+
+def test_service_controller_does_nothing_if_not_started(tmpdir: Any) -> None:
+    """Test that nothing happened if controller is not started."""
+    service_controller = get_system_controller(
+        single_instance=True, pid_file_path=Path(tmpdir) / "test.pid"
+    )
+
+    assert not service_controller.is_running()
+    service_controller.stop()
+    assert not service_controller.is_running()