diff --git a/libs/computer/computer/computer.py b/libs/computer/computer/computer.py index c25ad2bf..b77582ae 100644 --- a/libs/computer/computer/computer.py +++ b/libs/computer/computer/computer.py @@ -722,3 +722,173 @@ class Computer: tuple[float, float]: (x, y) coordinates in screenshot space """ return await self.interface.to_screenshot_coordinates(x, y) + + + # Add virtual environment management functions to computer interface + async def venv_install(self, venv_name: str, requirements: list[str]) -> tuple[str, str]: + """Install packages in a virtual environment. + + Args: + venv_name: Name of the virtual environment + requirements: List of package requirements to install + + Returns: + Tuple of (stdout, stderr) from the installation command + """ + requirements = requirements or [] + + # Create virtual environment if it doesn't exist + venv_path = f"~/.venvs/{venv_name}" + create_cmd = f"mkdir -p ~/.venvs && python3 -m venv {venv_path}" + + # Check if venv exists, if not create it + check_cmd = f"test -d {venv_path} || ({create_cmd})" + _, _ = await self.interface.run_command(check_cmd) + + # Install packages + requirements_str = " ".join(requirements) + install_cmd = f". {venv_path}/bin/activate && pip install {requirements_str}" + return await self.interface.run_command(install_cmd) + + async def venv_cmd(self, venv_name: str, command: str) -> tuple[str, str]: + """Execute a shell command in a virtual environment. + + Args: + venv_name: Name of the virtual environment + command: Shell command to execute in the virtual environment + + Returns: + Tuple of (stdout, stderr) from the command execution + """ + venv_path = f"~/.venvs/{venv_name}" + + # Check if virtual environment exists + check_cmd = f"test -d {venv_path}" + stdout, stderr = await self.interface.run_command(check_cmd) + + if stderr or "test:" in stdout: # venv doesn't exist + return "", f"Virtual environment '{venv_name}' does not exist. Create it first using venv_install." + + # Activate virtual environment and run command + full_command = f". {venv_path}/bin/activate && {command}" + return await self.interface.run_command(full_command) + + async def venv_exec(self, venv_name: str, python_func, *args, **kwargs): + """Execute Python function in a virtual environment using source code extraction. + + Args: + venv_name: Name of the virtual environment + python_func: A callable function to execute + *args: Positional arguments to pass to the function + **kwargs: Keyword arguments to pass to the function + + Returns: + The result of the function execution, or raises any exception that occurred + """ + import base64 + import inspect + import json + import textwrap + + try: + # Get function source code using inspect.getsource + source = inspect.getsource(python_func) + # Remove common leading whitespace (dedent) + func_source = textwrap.dedent(source).strip() + + # Get function name for execution + func_name = python_func.__name__ + + # Serialize args and kwargs as JSON (safer than dill for cross-version compatibility) + args_json = json.dumps(args, default=str) + kwargs_json = json.dumps(kwargs, default=str) + + except OSError as e: + raise Exception(f"Cannot retrieve source code for function {python_func.__name__}: {e}") + except Exception as e: + raise Exception(f"Failed to reconstruct function source: {e}") + + # Create Python code that will define and execute the function + python_code = f''' +import json +import traceback + +try: + # Define the function from source +{textwrap.indent(func_source, " ")} + + # Deserialize args and kwargs from JSON + args_json = """{args_json}""" + kwargs_json = """{kwargs_json}""" + args = json.loads(args_json) + kwargs = json.loads(kwargs_json) + + # Execute the function + result = {func_name}(*args, **kwargs) + + # Create success output payload + output_payload = {{ + "success": True, + "result": result, + "error": None + }} + +except Exception as e: + # Create error output payload + output_payload = {{ + "success": False, + "result": None, + "error": {{ + "type": type(e).__name__, + "message": str(e), + "traceback": traceback.format_exc() + }} + }} + +# Serialize the output payload as JSON +import json +output_json = json.dumps(output_payload, default=str) + +# Print the JSON output with markers +print(f"<<>>{{output_json}}<<>>") +''' + + # Encode the Python code in base64 to avoid shell escaping issues + encoded_code = base64.b64encode(python_code.encode('utf-8')).decode('ascii') + + # Execute the Python code in the virtual environment + python_command = f"python -c \"import base64; exec(base64.b64decode('{encoded_code}').decode('utf-8'))\"" + stdout, stderr = await self.venv_cmd(venv_name, python_command) + + # Parse the output to extract the payload + start_marker = "<<>>" + end_marker = "<<>>" + + # Print original stdout + print(stdout[:stdout.find(start_marker)]) + + if start_marker in stdout and end_marker in stdout: + start_idx = stdout.find(start_marker) + len(start_marker) + end_idx = stdout.find(end_marker) + + if start_idx < end_idx: + output_json = stdout[start_idx:end_idx] + + try: + # Decode and deserialize the output payload from JSON + output_payload = json.loads(output_json) + except Exception as e: + raise Exception(f"Failed to decode output payload: {e}") + + if output_payload["success"]: + return output_payload["result"] + else: + # Recreate and raise the original exception + error_info = output_payload["error"] + error_class = eval(error_info["type"]) + raise error_class(error_info["message"]) + else: + raise Exception("Invalid output format: markers found but no content between them") + else: + # Fallback: return stdout/stderr if no payload markers found + raise Exception(f"No output payload found. stdout: {stdout}, stderr: {stderr}") diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 00000000..998cbeaf --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +asyncio_mode = auto +markers = + asyncio: asyncio mark \ No newline at end of file diff --git a/tests/venv.py b/tests/venv.py new file mode 100644 index 00000000..4f9e3206 --- /dev/null +++ b/tests/venv.py @@ -0,0 +1,151 @@ +""" +Virtual Environment Testing Module +This module tests the ability to execute python code in a virtual environment within C/ua Containers. + +Required environment variables: +- CUA_API_KEY: API key for C/ua cloud provider +- CUA_CONTAINER_NAME: Name of the container to use +""" + +import os +import asyncio +import pytest +from pathlib import Path +import sys +import traceback + +# Load environment variables from .env file +project_root = Path(__file__).parent.parent +env_file = project_root / ".env" +print(f"Loading environment from: {env_file}") +from dotenv import load_dotenv + +load_dotenv(env_file) + +# Add paths to sys.path if needed +pythonpath = os.environ.get("PYTHONPATH", "") +for path in pythonpath.split(":"): + if path and path not in sys.path: + sys.path.insert(0, path) # Insert at beginning to prioritize + print(f"Added to sys.path: {path}") + +from computer.computer import Computer +from computer.providers.base import VMProviderType + + +@pytest.fixture(scope="session") +async def computer(): + """Shared Computer instance for all test cases.""" + # Create a remote Linux computer with C/ua + computer = Computer( + os_type="linux", + api_key=os.getenv("CUA_API_KEY"), + name=str(os.getenv("CUA_CONTAINER_NAME")), + provider_type=VMProviderType.CLOUD, + ) + + try: + await computer.run() + yield computer + finally: + await computer.stop() + + +# Sample test cases +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_install(computer): + """Test virtual environment creation and package installation.""" + # Create a test virtual environment and install requests + stdout, _ = await computer.venv_install("test_env", ["requests"]) + + # Check that installation was successful (no major errors) + assert "Successfully installed" in stdout or "Requirement already satisfied" in stdout + +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_cmd(computer): + """Test executing shell commands in virtual environment.""" + # Test Python version check + stdout, _ = await computer.venv_cmd("test_env", "python --version") + + assert "Python" in stdout + +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_exec(computer): + """Test executing Python functions in virtual environment.""" + def test_function(message="Hello World"): + import sys + return f"Python {sys.version_info.major}.{sys.version_info.minor}: {message}" + + result = await computer.venv_exec("test_env", test_function, message="Test successful!") + + assert "Python" in result + assert "Test successful!" in result + +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_exec_with_package(computer): + """Test executing Python functions that use installed packages.""" + def test_requests(): + import requests + return f"requests version: {requests.__version__}" + + result = await computer.venv_exec("test_env", test_requests) + + assert "requests version:" in result + +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_exec_error_handling(computer): + """Test error handling in venv_exec.""" + def test_error(): + raise ValueError("This is a test error") + + with pytest.raises(ValueError, match="This is a test error"): + await computer.venv_exec("test_env", test_error) + +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_exec_with_args_kwargs(computer): + """Test executing Python functions with args and kwargs that return an object.""" + def create_data_object(name, age, *hobbies, **metadata): + return { + "name": name, + "age": age, + "hobbies": list(hobbies), + "metadata": metadata, + "status": "active" + } + + args = ["Alice", 25, "reading", "coding"] + kwargs = {"location": "New York", "department": "Engineering"} + + result = await computer.venv_exec( + "test_env", + create_data_object, + *args, + **kwargs + ) + + assert result["name"] == "Alice" + assert result["age"] == 25 + assert result["hobbies"] == ["reading", "coding"] + assert result["metadata"]["location"] == "New York" + assert result["status"] == "active" + +@pytest.mark.asyncio(loop_scope="session") +async def test_venv_exec_stdout_capture(computer, capfd): + """Test capturing stdout from Python functions executed in virtual environment.""" + def hello_world_function(): + print("Hello World!") + return "Function completed" + + # Execute the function in the virtual environment + result = await computer.venv_exec("test_env", hello_world_function) + + # Capture stdout and stderr + out, _ = capfd.readouterr() + + # Assert the stdout contains our expected output + assert out == "Hello World!\n\n" + assert result == "Function completed" + +if __name__ == "__main__": + # Run tests directly + pytest.main([__file__, "-v"])