178 lines
5.2 KiB
Python
178 lines
5.2 KiB
Python
from http import HTTPStatus
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from e2b import AsyncVolume
|
|
from e2b.exceptions import NotFoundException
|
|
from e2b.api.client.models.volume_and_token import VolumeAndToken
|
|
from e2b.api.client.types import Response
|
|
import e2b.api.client.api.volumes.post_volumes as post_volumes_mod
|
|
import e2b.api.client.api.volumes.get_volumes as get_volumes_mod
|
|
import e2b.api.client.api.volumes.get_volumes_volume_id as get_volume_mod
|
|
import e2b.api.client.api.volumes.delete_volumes_volume_id as delete_volume_mod
|
|
|
|
# In-memory store for mock volumes
|
|
_volumes: dict[str, VolumeAndToken] = {}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_volume_api(monkeypatch, test_api_key):
|
|
monkeypatch.setenv("E2B_API_KEY", test_api_key)
|
|
_volumes.clear()
|
|
|
|
async def mock_post_volumes(*, client, body):
|
|
vol_id = str(uuid4())
|
|
token = f"vol-token-{uuid4()}"
|
|
vol = VolumeAndToken(volume_id=vol_id, name=body.name, token=token)
|
|
_volumes[vol_id] = vol
|
|
return Response(
|
|
status_code=HTTPStatus(201),
|
|
content=b"",
|
|
headers={},
|
|
parsed=vol,
|
|
)
|
|
|
|
async def mock_get_volumes(*, client):
|
|
return Response(
|
|
status_code=HTTPStatus(200),
|
|
content=b"",
|
|
headers={},
|
|
parsed=list(_volumes.values()),
|
|
)
|
|
|
|
async def mock_get_volume(volume_id, *, client):
|
|
vol = _volumes.get(volume_id)
|
|
if vol is None:
|
|
return Response(
|
|
status_code=HTTPStatus(404),
|
|
content=b"",
|
|
headers={},
|
|
parsed=None,
|
|
)
|
|
return Response(
|
|
status_code=HTTPStatus(200),
|
|
content=b"",
|
|
headers={},
|
|
parsed=vol,
|
|
)
|
|
|
|
async def mock_delete_volume(volume_id, *, client):
|
|
if volume_id not in _volumes:
|
|
return Response(
|
|
status_code=HTTPStatus(404),
|
|
content=b"",
|
|
headers={},
|
|
parsed=None,
|
|
)
|
|
del _volumes[volume_id]
|
|
return Response(
|
|
status_code=HTTPStatus(204),
|
|
content=b"",
|
|
headers={},
|
|
parsed=None,
|
|
)
|
|
|
|
monkeypatch.setattr(post_volumes_mod, "asyncio_detailed", mock_post_volumes)
|
|
monkeypatch.setattr(get_volumes_mod, "asyncio_detailed", mock_get_volumes)
|
|
monkeypatch.setattr(get_volume_mod, "asyncio_detailed", mock_get_volume)
|
|
monkeypatch.setattr(delete_volume_mod, "asyncio_detailed", mock_delete_volume)
|
|
|
|
|
|
async def test_create_volume():
|
|
vol = await AsyncVolume.create("test-volume")
|
|
|
|
assert vol is not None
|
|
assert vol.volume_id is not None
|
|
assert vol.name == "test-volume"
|
|
assert vol.token is not None
|
|
|
|
|
|
async def test_get_volume_info():
|
|
created = await AsyncVolume.create("info-volume")
|
|
info = await AsyncVolume.get_info(created.volume_id)
|
|
|
|
assert info.volume_id == created.volume_id
|
|
assert info.name == "info-volume"
|
|
assert info.token is not None
|
|
|
|
|
|
async def test_list_volumes():
|
|
await AsyncVolume.create("vol-a")
|
|
await AsyncVolume.create("vol-b")
|
|
|
|
volumes = await AsyncVolume.list()
|
|
|
|
assert len(volumes) == 2
|
|
names = sorted([v.name for v in volumes])
|
|
assert names == ["vol-a", "vol-b"]
|
|
|
|
|
|
async def test_list_volumes_empty():
|
|
volumes = await AsyncVolume.list()
|
|
assert len(volumes) == 0
|
|
|
|
|
|
async def test_destroy_volume():
|
|
vol = await AsyncVolume.create("to-delete")
|
|
result = await AsyncVolume.destroy(vol.volume_id)
|
|
|
|
assert result is True
|
|
|
|
volumes = await AsyncVolume.list()
|
|
assert len(volumes) == 0
|
|
|
|
|
|
async def test_destroy_nonexistent_volume():
|
|
result = await AsyncVolume.destroy("non-existent-id")
|
|
assert result is False
|
|
|
|
|
|
async def test_get_info_nonexistent_volume():
|
|
with pytest.raises(NotFoundException):
|
|
await AsyncVolume.get_info("non-existent-id")
|
|
|
|
|
|
async def test_create_volume_keeps_proxy_for_content_calls():
|
|
vol = await AsyncVolume.create(
|
|
"proxy-volume", proxy="http://user:pass@127.0.0.1:8080"
|
|
)
|
|
|
|
# The proxy is stored on the instance...
|
|
assert vol._proxy == "http://user:pass@127.0.0.1:8080"
|
|
|
|
# ...and instance methods (which build a VolumeConnectionConfig with no
|
|
# per-call proxy) pick it up rather than falling back to no proxy.
|
|
config = vol._get_volume_config()
|
|
assert config.proxy == "http://user:pass@127.0.0.1:8080"
|
|
|
|
|
|
async def test_volume_per_call_proxy_overrides_instance():
|
|
vol = await AsyncVolume.create("proxy-volume", proxy="http://127.0.0.1:8080")
|
|
|
|
config = vol._get_volume_config(proxy="http://127.0.0.1:9090")
|
|
assert config.proxy == "http://127.0.0.1:9090"
|
|
|
|
|
|
async def test_volume_full_lifecycle():
|
|
# Create
|
|
vol = await AsyncVolume.create("lifecycle-vol")
|
|
assert vol.name == "lifecycle-vol"
|
|
|
|
# Get info
|
|
info = await AsyncVolume.get_info(vol.volume_id)
|
|
assert info.name == "lifecycle-vol"
|
|
|
|
# List
|
|
volumes = await AsyncVolume.list()
|
|
assert len(volumes) == 1
|
|
assert volumes[0].volume_id == vol.volume_id
|
|
|
|
# Destroy
|
|
destroyed = await AsyncVolume.destroy(vol.volume_id)
|
|
assert destroyed is True
|
|
|
|
# List again
|
|
volumes = await AsyncVolume.list()
|
|
assert len(volumes) == 0
|