chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,280 @@
|
||||
import asyncio
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from e2b import (
|
||||
AsyncCommandHandle,
|
||||
AsyncSandbox,
|
||||
AsyncTemplate,
|
||||
AsyncVolume,
|
||||
CommandExitException,
|
||||
CommandHandle,
|
||||
LogEntry,
|
||||
Sandbox,
|
||||
Template,
|
||||
TemplateClass,
|
||||
Volume,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_api_key() -> str:
|
||||
"""Placeholder API key with a valid format for tests that don't hit the API."""
|
||||
return "e2b_" + "0" * 40
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
|
||||
def pytest_runtest_makereport(item, call):
|
||||
outcome = yield
|
||||
rep = outcome.get_result()
|
||||
if rep.when == "call":
|
||||
item._test_failed = rep.failed
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def sandbox_test_id():
|
||||
return f"test_{_generate_random_string()}"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def template():
|
||||
return "base"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def sandbox_factory(request, template, sandbox_test_id):
|
||||
def factory(*, template_name: str = template, **kwargs):
|
||||
metadata = kwargs.setdefault("metadata", dict())
|
||||
metadata.setdefault("sandbox_test_id", sandbox_test_id)
|
||||
|
||||
sandbox = Sandbox.create(template_name, **kwargs)
|
||||
|
||||
def finalizer():
|
||||
if getattr(request.node, "_test_failed", False):
|
||||
print(f"\n[TEST FAILED] Sandbox ID: {sandbox.sandbox_id}")
|
||||
sandbox.kill()
|
||||
|
||||
request.addfinalizer(finalizer)
|
||||
|
||||
return sandbox
|
||||
|
||||
return factory
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def sandbox(sandbox_factory):
|
||||
return sandbox_factory()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def async_sandbox_factory(request, template, sandbox_test_id):
|
||||
sandboxes: list = []
|
||||
|
||||
async def factory(*, template_name: str = template, **kwargs):
|
||||
metadata = kwargs.setdefault("metadata", dict())
|
||||
metadata.setdefault("sandbox_test_id", sandbox_test_id)
|
||||
|
||||
sandbox = await AsyncSandbox.create(template_name, **kwargs)
|
||||
sandboxes.append(sandbox)
|
||||
return sandbox
|
||||
|
||||
yield factory
|
||||
|
||||
if getattr(request.node, "_test_failed", False):
|
||||
for sandbox in sandboxes:
|
||||
print(f"\n[TEST FAILED] Sandbox ID: {sandbox.sandbox_id}")
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(sandbox.kill() for sandbox in sandboxes), return_exceptions=True
|
||||
)
|
||||
for sandbox, result in zip(sandboxes, results):
|
||||
if isinstance(result, BaseException):
|
||||
print(f"\n[TEARDOWN FAILED] Sandbox ID: {sandbox.sandbox_id}: {result!r}")
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def async_sandbox(async_sandbox_factory):
|
||||
return await async_sandbox_factory()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def build():
|
||||
def _build(
|
||||
template: TemplateClass,
|
||||
name: Optional[str] = None,
|
||||
skip_cache: bool = False,
|
||||
on_build_logs: Optional[Callable[[LogEntry], None]] = None,
|
||||
):
|
||||
build_name = name or f"e2b-test-{_generate_random_string()}"
|
||||
build_info: Dict[str, Optional[str]] = {"template_id": None, "build_id": None}
|
||||
|
||||
def capture_logs(log: LogEntry):
|
||||
import re
|
||||
|
||||
if "Template created with ID:" in log.message:
|
||||
match = re.search(
|
||||
r"Template created with ID: ([^,]+), Build ID: (.+)", log.message
|
||||
)
|
||||
if match:
|
||||
build_info["template_id"] = match.group(1)
|
||||
build_info["build_id"] = match.group(2)
|
||||
if on_build_logs:
|
||||
on_build_logs(log)
|
||||
|
||||
try:
|
||||
return Template.build(
|
||||
template,
|
||||
build_name,
|
||||
cpu_count=1,
|
||||
memory_mb=1024,
|
||||
skip_cache=skip_cache,
|
||||
on_build_logs=capture_logs,
|
||||
)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"\n[BUILD FAILED] name={build_name}, "
|
||||
f"template_id={build_info['template_id']}, "
|
||||
f"build_id={build_info['build_id']}, error={e}"
|
||||
)
|
||||
raise
|
||||
|
||||
return _build
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def async_build():
|
||||
async def _async_build(
|
||||
template: TemplateClass,
|
||||
name: Optional[str] = None,
|
||||
skip_cache: bool = False,
|
||||
on_build_logs: Optional[Callable[[LogEntry], None]] = None,
|
||||
):
|
||||
build_name = name or f"e2b-test-{_generate_random_string()}"
|
||||
build_info: Dict[str, Optional[str]] = {"template_id": None, "build_id": None}
|
||||
|
||||
def capture_logs(log: LogEntry):
|
||||
import re
|
||||
|
||||
if "Template created with ID:" in log.message:
|
||||
match = re.search(
|
||||
r"Template created with ID: ([^,]+), Build ID: (.+)", log.message
|
||||
)
|
||||
if match:
|
||||
build_info["template_id"] = match.group(1)
|
||||
build_info["build_id"] = match.group(2)
|
||||
if on_build_logs:
|
||||
on_build_logs(log)
|
||||
|
||||
try:
|
||||
return await AsyncTemplate.build(
|
||||
template,
|
||||
build_name,
|
||||
cpu_count=1,
|
||||
memory_mb=1024,
|
||||
skip_cache=skip_cache,
|
||||
on_build_logs=capture_logs,
|
||||
)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"\n[BUILD FAILED] name={build_name}, "
|
||||
f"template_id={build_info['template_id']}, "
|
||||
f"build_id={build_info['build_id']}, error={e}"
|
||||
)
|
||||
raise
|
||||
|
||||
return _async_build
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def debug():
|
||||
return os.getenv("E2B_DEBUG") is not None
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def skip_by_debug(request, debug):
|
||||
if request.node.get_closest_marker("skip_debug"):
|
||||
if debug:
|
||||
pytest.skip("skipped because E2B_DEBUG is set")
|
||||
|
||||
|
||||
class Helpers:
|
||||
@staticmethod
|
||||
def catch_cmd_exit_error_in_background(cmd: AsyncCommandHandle):
|
||||
disabled = False
|
||||
|
||||
async def wait_for_exit():
|
||||
try:
|
||||
await cmd.wait()
|
||||
except CommandExitException as e:
|
||||
if not disabled:
|
||||
assert False, (
|
||||
f"command failed with exit code {e.exit_code}: {e.stderr}"
|
||||
)
|
||||
|
||||
asyncio.create_task(wait_for_exit())
|
||||
|
||||
def disable():
|
||||
nonlocal disabled
|
||||
disabled = True
|
||||
|
||||
return disable
|
||||
|
||||
@staticmethod
|
||||
def check_cmd_exit_error(cmd: CommandHandle):
|
||||
try:
|
||||
cmd.wait()
|
||||
except CommandExitException as e:
|
||||
assert False, f"command failed with exit code {e.exit_code}: {e.stderr}"
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def helpers():
|
||||
return Helpers
|
||||
|
||||
|
||||
def _generate_random_string(length: int = 8) -> str:
|
||||
return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))
|
||||
|
||||
|
||||
def _skip_unless_volume_tests_enabled():
|
||||
if os.getenv("ENABLE_VOLUME_TESTS") is None:
|
||||
pytest.skip("skipped because ENABLE_VOLUME_TESTS is not set")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def volume(request):
|
||||
_skip_unless_volume_tests_enabled()
|
||||
vol = Volume.create(f"test-vol-{_generate_random_string()}")
|
||||
|
||||
def finalizer():
|
||||
if getattr(request.node, "_test_failed", False):
|
||||
print(f"\n[TEST FAILED] Volume ID: {vol.volume_id}")
|
||||
try:
|
||||
Volume.destroy(vol.volume_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
request.addfinalizer(finalizer)
|
||||
return vol
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def async_volume(request):
|
||||
_skip_unless_volume_tests_enabled()
|
||||
vol = await AsyncVolume.create(f"test-vol-{_generate_random_string()}")
|
||||
try:
|
||||
yield vol
|
||||
finally:
|
||||
if getattr(request.node, "_test_failed", False):
|
||||
print(f"\n[TEST FAILED] Volume ID: {vol.volume_id}")
|
||||
try:
|
||||
await AsyncVolume.destroy(vol.volume_id)
|
||||
except Exception:
|
||||
pass
|
||||
Reference in New Issue
Block a user