forked from eval-protocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_agent_resources.py
More file actions
426 lines (386 loc) · 18.7 KB
/
Copy pathtest_agent_resources.py
File metadata and controls
426 lines (386 loc) · 18.7 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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
"""
Tests for V2 ForkableResource implementations.
"""
import asyncio # For async tests
import copy
import pickle
import sqlite3
import tarfile
import tempfile
from pathlib import Path
from typing import Any, Dict # Added Dict, Any
from unittest.mock import AsyncMock # For mocking async methods if needed later
import pytest
import pytest_asyncio # Import pytest_asyncio
from eval_protocol.agent.resources.docker_resource import (
DOCKER_DAEMON_AVAILABLE,
DOCKER_SDK_AVAILABLE,
DockerResource,
)
from eval_protocol.agent.resources.filesystem_resource import FileSystemResource
from eval_protocol.agent.resources.python_state_resource import PythonStateResource
from eval_protocol.agent.resources.sql_resource import SQLResource
@pytest.mark.asyncio
class TestPythonStateResource:
"""Tests for the PythonStateResource."""
async def test_setup_initial_state(self):
config: Dict[str, Any] = {"initial_state": {"key1": "value1", "count": 10}}
resource = PythonStateResource()
await resource.setup(config)
observation = await resource.get_observation()
assert observation == {"key1": "value1", "count": 10}
assert resource._state is not config["initial_state"]
async def test_setup_empty_config(self):
config: Dict[str, Any] = {}
resource = PythonStateResource()
await resource.setup(config)
assert await resource.get_observation() == {}
async def test_fork_state_copy(self):
initial_state = {"data": [1, 2], "nested": {"a": "b"}}
config = {"initial_state": initial_state}
original_resource = PythonStateResource()
await original_resource.setup(config)
forked_resource = await original_resource.fork()
assert await forked_resource.get_observation() == initial_state
assert forked_resource._state is not original_resource._state
assert forked_resource._state["data"] is not original_resource._state["data"]
assert forked_resource._state["nested"] is not original_resource._state["nested"]
async def test_fork_independent_state(self):
config = {"initial_state": {"key": "original_value", "num": 1}}
original_resource = PythonStateResource()
await original_resource.setup(config)
forked_resource = await original_resource.fork()
await forked_resource.step("update_state", {"key": "forked_value"})
await forked_resource.step("update_state", {"num": 2})
assert (await original_resource.get_observation())["key"] == "original_value"
assert (await original_resource.get_observation())["num"] == 1
assert (await forked_resource.get_observation())["key"] == "forked_value"
assert (await forked_resource.get_observation())["num"] == 2
await original_resource.step("update_state", {"key": "new_original_value"})
await original_resource.step("update_state", {"num": 0})
assert (await forked_resource.get_observation())["key"] == "forked_value"
assert (await original_resource.get_observation())["key"] == "new_original_value"
async def test_checkpoint_and_restore(self):
initial_config = {"initial_state": {"val": 42, "items": ["a", "b"]}}
resource1 = PythonStateResource()
await resource1.setup(initial_config)
await resource1.step("update_state", {"val": 43, "new_item": "c"})
checkpoint_data = await resource1.checkpoint()
assert isinstance(checkpoint_data, bytes)
resource2 = PythonStateResource()
await resource2.setup({"initial_state": {"val": 0, "items": []}})
await resource2.restore(checkpoint_data)
expected_state = {"val": 43, "items": ["a", "b"], "new_item": "c"}
assert await resource2.get_observation() == expected_state
assert await resource1.get_observation() == expected_state
with pytest.raises((pickle.UnpicklingError, TypeError, AttributeError, ValueError)):
await resource2.restore(b"invalid pickle data")
async def test_step_update_state(self):
resource = PythonStateResource()
await resource.setup({"initial_state": {"a": 1}})
result = await resource.step("update_state", {"b": 2, "a": 3})
expected_state = {"a": 3, "b": 2}
assert await resource.get_observation() == expected_state
assert result == expected_state
async def test_step_get_value(self):
resource = PythonStateResource()
await resource.setup({"initial_state": {"name": "test", "value": 100}})
assert await resource.step("get_value", {"key": "name"}) == "test"
assert await resource.step("get_value", {"key": "value"}) == 100
assert await resource.step("get_value", {"key": "nonexistent"}) is None
with pytest.raises(ValueError, match="Missing 'key'"):
await resource.step("get_value", {})
async def test_step_unknown_action(self):
resource = PythonStateResource()
await resource.setup({})
with pytest.raises(NotImplementedError):
await resource.step("non_existent_action", {})
async def test_get_observation_is_copy(self):
resource = PythonStateResource()
initial_state = {"list_key": [1, 2]}
await resource.setup({"initial_state": initial_state})
obs1 = await resource.get_observation()
obs1["list_key"].append(3)
assert (await resource.get_observation())["list_key"] == [1, 2]
async def test_get_tools_spec(self):
resource = PythonStateResource()
specs = await resource.get_tools_spec()
assert isinstance(specs, list)
assert len(specs) == 2
tool_names = [spec["function"]["name"] for spec in specs]
assert "update_state" in tool_names
assert "get_value" in tool_names
async def test_close_clears_state(self):
resource = PythonStateResource()
await resource.setup({"initial_state": {"a": 1}})
assert await resource.get_observation() == {"a": 1}
await resource.close()
assert await resource.get_observation() == {}
assert resource._config == {}
assert resource._state == {}
@pytest.mark.asyncio
class TestSQLResource:
"""Tests for the SQLResource."""
async def test_setup_sqlite_with_schema_file_and_seed_file(self, tmp_path):
schema_content = "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT);"
seed_content = "INSERT INTO users (name) VALUES ('Alice'); INSERT INTO users (name) VALUES ('Bob');"
schema_file = tmp_path / "schema.sql"
schema_file.write_text(schema_content)
seed_file = tmp_path / "seed.sql"
seed_file.write_text(seed_content)
config = {
"db_type": "sqlite",
"db_name": "test_setup.sqlite",
"schema_file": str(schema_file),
"seed_data_file": str(seed_file),
}
resource = SQLResource()
try:
await resource.setup(config)
assert resource._db_path is not None and resource._db_path.exists()
conn = sqlite3.connect(resource._db_path)
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM users")
assert cursor.fetchone()[0] == 2
conn.close()
finally:
await resource.close()
async def test_fork_sqlite(self):
config = {
"db_type": "sqlite",
"db_name": "original.sqlite",
"schema_sql": "CREATE TABLE data (val TEXT);",
}
original_resource = SQLResource()
try:
await original_resource.setup(config)
await original_resource.step("execute_sql", {"query": "INSERT INTO data (val) VALUES ('initial')"})
forked_resource = await original_resource.fork()
try:
assert forked_resource._db_path is not None and forked_resource._db_path.exists()
res_fork = await forked_resource.step(
"execute_sql",
{"query": "SELECT val FROM data", "fetch_mode": "one"},
)
assert res_fork["val"] == "initial"
await forked_resource.step(
"execute_sql",
{"query": "UPDATE data SET val = 'forked_changed' WHERE val = 'initial'"},
)
res_orig = await original_resource.step(
"execute_sql",
{"query": "SELECT val FROM data", "fetch_mode": "one"},
)
assert res_orig["val"] == "initial"
finally:
await forked_resource.close()
finally:
await original_resource.close()
async def test_checkpoint_and_restore_sqlite(self):
config = {
"db_type": "sqlite",
"db_name": "chkpt_orig.sqlite",
"schema_sql": "CREATE TABLE log (msg TEXT);",
}
res1 = SQLResource()
try:
await res1.setup(config)
await res1.step("execute_sql", {"query": "INSERT INTO log (msg) VALUES ('message1')"})
checkpoint_info = await res1.checkpoint()
checkpoint_file = Path(checkpoint_info["checkpoint_path"])
assert checkpoint_file.exists()
res2 = SQLResource()
try:
await res2.setup({"db_type": "sqlite", "db_name": "chkpt_res2_init.sqlite"})
await res2.restore(checkpoint_info)
count_res2 = await res2.step(
"execute_sql",
{"query": "SELECT COUNT(*) FROM log", "fetch_mode": "val"},
)
assert count_res2 == 1
finally:
await res2.close()
if checkpoint_file.exists():
checkpoint_file.unlink()
finally:
await res1.close()
async def test_close_sqlite(self):
config = {"db_type": "sqlite", "db_name": "closetest.sqlite"}
resource = SQLResource()
await resource.setup(config)
db_path = resource._db_path
assert db_path is not None and db_path.exists()
await resource.close()
assert not db_path.exists()
@pytest.mark.asyncio
class TestFileSystemResource:
"""Tests for the FileSystemResource."""
async def test_setup_with_initial_files(self):
config = {
"initial_files": {
"file1.txt": "Hello World",
"subdir/file2.py": "print('Python code')",
}
}
resource = FileSystemResource()
try:
await resource.setup(config)
assert resource._managed_dir_path is not None and resource._managed_dir_path.exists()
# Add assertion for mypy
assert resource._managed_dir_path is not None
assert (resource._managed_dir_path / "file1.txt").read_text() == "Hello World"
assert (resource._managed_dir_path / "subdir" / "file2.py").read_text() == "print('Python code')"
finally:
await resource.close()
async def test_fork_filesystem(self):
config = {"initial_files": {"original.txt": "original content"}}
original_resource = FileSystemResource()
try:
await original_resource.setup(config)
forked_resource = await original_resource.fork()
try:
assert forked_resource._managed_dir_path is not None and forked_resource._managed_dir_path.exists()
# Add assertion for mypy
assert forked_resource._managed_dir_path is not None
assert (forked_resource._managed_dir_path / "original.txt").read_text() == "original content"
(forked_resource._managed_dir_path / "forked_new.txt").write_text("forked specific")
# Add assertion for mypy
assert original_resource._managed_dir_path is not None
assert not (original_resource._managed_dir_path / "forked_new.txt").exists()
finally:
await forked_resource.close()
finally:
await original_resource.close()
async def test_checkpoint_and_restore_filesystem(self, tmp_path):
config = {"initial_files": {"data.txt": "checkpoint me"}}
res1 = FileSystemResource()
try:
await res1.setup(config)
checkpoint_info = await res1.checkpoint()
checkpoint_archive = Path(checkpoint_info["checkpoint_path"])
assert checkpoint_archive.exists()
res2 = FileSystemResource()
try:
await res2.setup({"base_dir_name": "fs_res2_init"})
await res2.restore(checkpoint_info)
assert res2._managed_dir_path is not None # Ensure for mypy
assert (res2._managed_dir_path / "data.txt").read_text() == "checkpoint me"
finally:
await res2.close()
if checkpoint_archive.exists():
checkpoint_archive.unlink()
finally:
await res1.close()
async def test_close_filesystem(self):
resource = FileSystemResource()
await resource.setup({"initial_files": {"a.txt": "content"}})
managed_path = resource._managed_dir_path
assert managed_path is not None and managed_path.exists()
await resource.close()
assert not managed_path.exists()
pytestmark_docker = pytest.mark.skipif(
not DOCKER_SDK_AVAILABLE or not DOCKER_DAEMON_AVAILABLE,
reason="Docker SDK not installed or Docker daemon not running",
)
@pytestmark_docker
@pytest.mark.asyncio
class TestDockerResource:
"""Tests for the DockerResource."""
TEST_IMAGE = "alpine:latest"
@pytest_asyncio.fixture(scope="function") # Changed to pytest_asyncio.fixture
async def docker_resource(self): # Made async fixture
resource = DockerResource()
try:
client = resource._client
await asyncio.to_thread(client.images.get, self.TEST_IMAGE)
except Exception:
try:
print(f"\nPulling Docker test image: {self.TEST_IMAGE}...")
await asyncio.to_thread(client.images.pull, self.TEST_IMAGE)
print("Pull complete.")
except Exception as e:
pytest.skip(f"Failed to pull Docker image {self.TEST_IMAGE}: {e}")
yield resource
await resource.close()
async def test_setup_docker_container(self, docker_resource: DockerResource):
config = {
"image_name": self.TEST_IMAGE,
"docker_run_options": {"command": ["tail", "-f", "/dev/null"]},
}
await docker_resource.setup(config)
assert docker_resource._container is not None
# Docker SDK calls are blocking, wrap in to_thread for async tests
container_status = await asyncio.to_thread(getattr, docker_resource._container, "status")
assert container_status == "running"
result = await docker_resource.step("exec_command", {"command": "echo hello"})
assert result["exit_code"] == 0 and "hello" in result["output"]
async def test_fork_docker_container(self, docker_resource: DockerResource):
await docker_resource.setup(
{
"image_name": self.TEST_IMAGE,
"docker_run_options": {"command": ["tail", "-f", "/dev/null"]},
}
)
await docker_resource.step("exec_command", {"command": "touch /tmp/original_file.txt"})
forked_resource = await docker_resource.fork()
try:
assert forked_resource._container is not None
res_fork = await forked_resource.step("exec_command", {"command": "ls /tmp/original_file.txt"})
assert res_fork["exit_code"] == 0
finally:
await forked_resource.close()
async def test_checkpoint_and_restore_docker(self, docker_resource: DockerResource):
config = {
"image_name": self.TEST_IMAGE,
"docker_run_options": {"command": ["tail", "-f", "/dev/null"]},
}
await docker_resource.setup(config)
create_file_command = "sh -c \"echo 'initial_data' > /data.txt\""
create_file_result = await docker_resource.step("exec_command", {"command": create_file_command})
assert create_file_result["exit_code"] == 0, (
f"Failed to create /data.txt with '{create_file_command}': {create_file_result['output']}"
)
# Optionally, verify file content immediately after creation in the source container
verify_result = await docker_resource.step("exec_command", {"command": "cat /data.txt"})
assert verify_result["exit_code"] == 0, f"Failed to cat /data.txt after creation: {verify_result['output']}"
assert "initial_data" in verify_result["output"], (
f"/data.txt content mismatch after creation: {verify_result['output']}"
)
checkpoint_info = await docker_resource.checkpoint()
checkpoint_image_id = checkpoint_info["image_id"]
restored_resource = DockerResource()
try:
await restored_resource.setup(config)
await restored_resource.restore(checkpoint_info)
result = await restored_resource.step("exec_command", {"command": "cat /data.txt"})
assert "initial_data" in result["output"]
finally:
await restored_resource.close()
# Note: The original image_name from config is self.TEST_IMAGE.
# The committed image (checkpoint_image_id) will be different.
# We should always try to clean up the checkpoint_image_id if it exists.
if checkpoint_image_id:
try:
# Check if it's the base image to avoid removing it if tests are re-run without pulling
base_image_obj = await asyncio.to_thread(docker_resource._client.images.get, self.TEST_IMAGE)
if base_image_obj.id != checkpoint_image_id:
await asyncio.to_thread(
docker_resource._client.images.remove,
image=checkpoint_image_id,
force=True,
)
except Exception as e: # Catch NotFound from get or APIError from remove
print(f"Warning: failed to cleanup checkpoint image {checkpoint_image_id}: {e}")
async def test_close_docker_resource(self, docker_resource: DockerResource):
await docker_resource.setup(
{
"image_name": self.TEST_IMAGE,
"docker_run_options": {"command": ["tail", "-f", "/dev/null"]},
}
)
assert docker_resource._container is not None # Ensure container exists before accessing id
container_id = docker_resource._container.id
await docker_resource.close()
assert docker_resource._container is None
with pytest.raises(Exception):
await asyncio.to_thread(docker_resource._client.containers.get, container_id)