-
Notifications
You must be signed in to change notification settings - Fork 21
Expand file tree
/
Copy pathtest_e2b_integration.py
More file actions
executable file
·74 lines (57 loc) · 2.28 KB
/
Copy pathtest_e2b_integration.py
File metadata and controls
executable file
·74 lines (57 loc) · 2.28 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
#!/usr/bin/env python
"""
Test script to verify E2B integration works end-to-end.
This script runs a simple Python code example using the E2B code execution
reward function and verifies that it works correctly.
"""
import os
import pytest
from eval_protocol.models import EvaluateResult
from eval_protocol.rewards.code_execution import _HAS_E2B, e2b_code_execution_reward
from tests.conftest import skip_e2b
@skip_e2b
def test_e2b_python_integration():
"""Test that E2B integration works correctly for Python code."""
# Verify API key is available
api_key = os.environ.get("E2B_API_KEY")
if not api_key:
pytest.skip("E2B_API_KEY environment variable is not set")
# Simple test case with Python code
messages = [
{"role": "user", "content": "Write a Python function to add two numbers."},
{
"role": "assistant",
"content": """Here's a Python function to add two numbers:
```python
def add(a, b):
return a + b
print(add(2, 3))
```
This function takes two numbers as input and returns their sum.""",
},
]
expected_output = "5"
# Evaluate the code using E2B
result = e2b_code_execution_reward(
messages=messages,
expected_output=expected_output,
language="python",
api_key=api_key,
timeout=15, # Increase timeout for first-time sandbox creation
)
# Verify the result
assert isinstance(result, EvaluateResult)
# If we get a sandbox error, consider the test successful if it contains expected errors
if result.score == 0.0 and "execution_result" in result.metrics:
error_msg = result.metrics["execution_result"].reason
if "sandbox timeout" in error_msg or "sandbox was not found" in error_msg or "Invalid API key" in error_msg:
pytest.skip(f"Skipping due to E2B connection issue: {error_msg}")
# Also check for error in main result reason
if result.reason and "Invalid API key" in result.reason:
pytest.skip("Skipping due to invalid E2B API key")
# Otherwise, it should be a successful result
assert result.score == 1.0
assert "execution_result" in result.metrics
assert result.metrics["execution_result"].is_score_valid is True
if __name__ == "__main__":
pytest.main(["-xvs", __file__])