Use pytest fixtures to run tests on local env and agent env

python/tests/conftest.py

@@ -3,9 +3,18 @@
 import platform
 import sys
 from pathlib import Path
+from typing import TYPE_CHECKING, Generator
 
 import pytest
 
+from binharness.bootstrap.docker import bootstrap_env_from_image
+from binharness.bootstrap.subprocess import SubprocessAgent
+from binharness.localenvironment import LocalEnvironment
+
+if TYPE_CHECKING:
+    from binharness.agentenvironment import AgentEnvironment
+    from binharness.types.environment import Environment
+
 try:
     import docker
 except ImportError:
@@ -24,17 +33,15 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
 
     # Skip docker tests if docker is not installed
     if "docker" in list(item.iter_markers()):
-        try:
-            import docker
-        except ImportError:
+        if docker is None:
             pytest.skip("docker is not installed")
         try:
             docker.from_env()
         except docker.errors.DockerException:
             pytest.skip("docker is not running")
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
 def agent_binary_host() -> str:
     expected_path = (
         Path(__file__).parent.parent.parent / "target" / "debug" / "bh_agent_server"
@@ -44,7 +51,7 @@ def agent_binary_host() -> str:
     return str(expected_path)
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
 def agent_binary_linux_host_arch() -> str:
     arch = platform.machine()
     # Fixup arch for arm64
@@ -63,11 +70,52 @@ def agent_binary_linux_host_arch() -> str:
     return str(expected_path)
 
 
-@pytest.fixture()
+@pytest.fixture(scope="session")
+def local_env() -> LocalEnvironment:
+    return LocalEnvironment()
+
+
+@pytest.fixture(scope="session")
+@pytest.mark.linux()
+def linux_local_env() -> LocalEnvironment:
+    return LocalEnvironment()
+
+
+@pytest.fixture(scope="session")
+def agent_env(agent_binary_host: str) -> Generator[AgentEnvironment, None, None]:
+    agent = SubprocessAgent(Path(agent_binary_host))
+    yield agent.get_environment(0)
+    agent.stop()
+
+
+@pytest.fixture(scope="session")
+@pytest.mark.docker()
 def docker_client() -> docker.DockerClient:
-    if docker is None:
-        pytest.skip("docker is not installed")
     try:
-        return docker.from_env()
+        client = docker.from_env()
     except docker.errors.DockerException:
         pytest.skip("docker is not running")
+    yield client
+    client.close()
+
+
+@pytest.fixture(scope="session")
+def docker_env(
+    docker_client: docker.APIClient, agent_binary_linux_host_arch: str
+) -> Generator[AgentEnvironment, None, None]:
+    agent = bootstrap_env_from_image(
+        agent_binary_linux_host_arch, "ubuntu:22.04", docker_client=docker_client
+    )
+    yield agent.get_environment(0)
+    agent.container.stop()
+    agent.container.remove()
+
+
+@pytest.fixture(params=["local_env", "agent_env"], scope="session")
+def env(request) -> Environment:  # noqa: ANN001
+    return request.getfixturevalue(request.param)  # type: ignore[no-any-return]
+
+
+@pytest.fixture(params=["linux_local_env", "docker_env"], scope="session")
+def linux_env(request) -> Environment:  # noqa: ANN001
+    return request.getfixturevalue(request.param)  # type: ignore[no-any-return]

python/tests/test_busybox.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import pytest
 
 from binharness.common.busybox import BusyboxInjection
-from binharness.localenvironment import LocalEnvironment
+
+if TYPE_CHECKING:
+    from binharness import Environment
 
 
 @pytest.mark.linux()
-def test_busybox_injection() -> None:
-    env = LocalEnvironment()
+def test_busybox_injection(env: Environment) -> None:
     busybox_injection = BusyboxInjection()
     busybox_injection.install(env)
     assert busybox_injection.run("true").wait() == 0
@@ -26,19 +29,18 @@ def test_busybox_injection() -> None:
 
 
 @pytest.mark.linux()
-def test_busbox_injection_mktemp() -> None:
-    env = LocalEnvironment()
+def test_busbox_injection_mktemp(env: Environment) -> None:
     busybox_injection = BusyboxInjection()
     busybox_injection.install(env)
     assert busybox_injection.mktemp().is_file()
     assert busybox_injection.mktemp(directory=True).is_dir()
 
 
+# TODO: Make this work on agents, probably an issue in wait()
 @pytest.mark.linux()
-def test_nc_interaction() -> None:
-    env = LocalEnvironment()
+def test_nc_interaction(local_env: Environment) -> None:
     busybox = BusyboxInjection()
-    busybox.install(env)
+    busybox.install(local_env)
 
     server_proc = busybox.nc("localhost", 10001, listen=True)
     client_proc = busybox.nc("localhost", 10001)
@@ -54,3 +56,8 @@ def test_nc_interaction() -> None:
     assert server_proc.stdout.readline() == b"hello\n"
     assert client_proc.stdout is not None
     assert client_proc.stdout.readline() == b"hello back\n"
+
+    client_proc.stdin.close()
+    client_proc.wait()
+    server_proc.stdin.close()
+    server_proc.wait()

python/tests/test_localenvironment.py → python/tests/test_environment.py

@@ -2,23 +2,24 @@
 
 import pathlib
 import tempfile
+from typing import TYPE_CHECKING
 
 import pytest
 
 from binharness.common.busybox import BusyboxInjection
-from binharness.localenvironment import LocalEnvironment
 
+if TYPE_CHECKING:
+    from binharness import Environment
 
-def test_run_command() -> None:
-    env = LocalEnvironment()
+
+def test_run_command(env: Environment) -> None:
     proc = env.run_command(["echo", "hello"])
     stdout, _ = proc.communicate()
     assert proc.returncode == 0
     assert stdout == b"hello\n"
 
 
-def test_inject_files() -> None:
-    env = LocalEnvironment()
+def test_inject_files(env: Environment) -> None:
     env_temp = env.get_tempdir()
     with tempfile.TemporaryDirectory() as tmp_dir:
         tmp_path = pathlib.Path(tmp_dir)
@@ -32,14 +33,14 @@ def test_inject_files() -> None:
         assert local_file.read_text() == "hello"
 
 
-def test_get_tempdir() -> None:
-    env = LocalEnvironment()
-    assert env.get_tempdir() == pathlib.Path(tempfile.gettempdir())
+# TODO: Need to think about how to handle this test with remote environments
+def test_get_tempdir(local_env: Environment) -> None:
+    assert local_env.get_tempdir() == pathlib.Path(tempfile.gettempdir())
 
 
+# TODO: Need to think about how to test these with non-linux environments
 @pytest.mark.linux()
-def test_stdout() -> None:
-    env = LocalEnvironment()
+def test_stdout(env: Environment) -> None:
     busybox = BusyboxInjection()
     busybox.install(env)
     proc = busybox.shell("echo hello")
@@ -48,8 +49,7 @@ def test_stdout() -> None:
 
 
 @pytest.mark.linux()
-def test_stderr() -> None:
-    env = LocalEnvironment()
+def test_stderr(env: Environment) -> None:
     busybox = BusyboxInjection()
     busybox.install(env)
     proc = busybox.shell("echo hello 1>&2")
@@ -58,8 +58,7 @@ def test_stderr() -> None:
 
 
 @pytest.mark.linux()
-def test_process_poll() -> None:
-    env = LocalEnvironment()
+def test_process_poll(env: Environment) -> None:
     busybox = BusyboxInjection()
     busybox.install(env)
     proc = busybox.run("head")

python/tests/test_executor.py

@@ -4,24 +4,25 @@
 
 import pytest
 
-from binharness.common.busybox import BusyboxShellExecutor
-from binharness.localenvironment import LocalEnvironment
-from binharness.types.executor import (
+from binharness import (
+    Environment,
     ExecutorEnvironmentMismatchError,
+    InjectionNotInstalledError,
+    LocalEnvironment,
+    Target,
 )
-from binharness.types.injection import InjectionNotInstalledError
-from binharness.types.target import Target
+from binharness.common.busybox import BusyboxShellExecutor
 
 
 @pytest.mark.linux()
-def test_busybox_injection_without_install() -> None:
-    env = LocalEnvironment()
+def test_busybox_injection_without_install(env: Environment) -> None:
     target = Target(env, Path("/usr/bin/true"))
     busybox_shell = BusyboxShellExecutor()
     with pytest.raises(InjectionNotInstalledError):
         assert busybox_shell.run_target(target)
 
 
+# TODO: Maybe LocalEnvironment objects should be interchangable?
 @pytest.mark.linux()
 def test_busybox_injection_different_environment() -> None:
     env1 = LocalEnvironment()

python/tests/test_inject.py

@@ -2,30 +2,31 @@
 
 import stat
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 
-from binharness.localenvironment import LocalEnvironment
 from binharness.types.injection import (
     ExecutableInjection,
     Injection,
     InjectionAlreadyInstalledError,
     InjectionNotInstalledError,
 )
 
+if TYPE_CHECKING:
+    from binharness.types.environment import Environment
+
 
 @pytest.mark.linux()
-def test_inject_true() -> None:
-    env = LocalEnvironment()
+def test_inject_true(env: Environment) -> None:
     true_injection = Injection(Path("/usr/bin/true"))
     true_injection.install(env)
     assert true_injection.env_path is not None
     assert stat.S_ISREG(env.stat(true_injection.env_path).mode)
 
 
 @pytest.mark.linux()
-def test_inject_true_executable() -> None:
-    env = LocalEnvironment()
+def test_inject_true_executable(env: Environment) -> None:
     true_injection = ExecutableInjection(Path("/usr/bin/true"))
     true_injection.install(env)
     assert true_injection.env_path is not None
@@ -36,17 +37,15 @@ def test_inject_true_executable() -> None:
 
 
 @pytest.mark.linux()
-def test_inject_true_executable_twice() -> None:
-    env = LocalEnvironment()
+def test_inject_true_executable_twice(env: Environment) -> None:
     true_injection = ExecutableInjection(Path("/usr/bin/true"))
     true_injection.install(env)
     with pytest.raises(InjectionAlreadyInstalledError):
         true_injection.install(env)
 
 
 @pytest.mark.linux()
-def test_inject_two_true_executables() -> None:
-    env = LocalEnvironment()
+def test_inject_two_true_executables(env: Environment) -> None:
     true_injection_1 = ExecutableInjection(Path("/usr/bin/true"))
     true_injection_1.install(env)
     true_injection_2 = ExecutableInjection(Path("/usr/bin/true"))

python/tests/test_qemu_executor.py

@@ -4,22 +4,20 @@
 
 import pytest
 
-from binharness import LocalEnvironment, Target
+from binharness import Environment, Target
 from binharness.common.qemu import QemuExecutor
 
 
 @pytest.mark.linux()
-def test_run_true() -> None:
-    env = LocalEnvironment()
+def test_run_true(env: Environment) -> None:
     target = Target(env, Path("/bin/true"))
     qemu = QemuExecutor()
     qemu.install(env)
     assert qemu.run_target(target).wait() == 0
 
 
 @pytest.mark.linux()
-def test_run_strace() -> None:
-    env = LocalEnvironment()
+def test_run_strace(env: Environment) -> None:
     target = Target(env, Path("/bin/true"))
     qemu = QemuExecutor()
     qemu.install(env)

python/tests/test_target_local.py

@@ -4,24 +4,19 @@
 
 import pytest
 
-from binharness import LocalEnvironment
+from binharness import Environment, NullExecutor, Target
 from binharness.common.busybox import BusyboxShellExecutor
-from binharness.types.executor import NullExecutor
-from binharness.types.target import Target
 
 
-@pytest.mark.linux()
-def test_run_target() -> None:
-    env = LocalEnvironment()
+def test_run_target(env: Environment) -> None:
     target = Target(env, Path("true"))
     executor = NullExecutor()
     proc = executor.run_target(target)
     assert proc.wait() == 0
 
 
 @pytest.mark.linux()
-def test_run_target_busybox() -> None:
-    env = LocalEnvironment()
+def test_run_target_busybox(env: Environment) -> None:
     target = Target(env, Path("true"))
     executor = BusyboxShellExecutor()
     executor.install(env)