|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import pytest |
| 4 | +from pydantic import ValidationError |
3 | 5 |
|
4 | 6 | from together.resources.code_interpreter import CodeInterpreter |
5 | 7 | from together.together_response import TogetherResponse |
@@ -326,3 +328,127 @@ def test_code_interpreter_session_management(mocker): |
326 | 328 |
|
327 | 329 | # Second call should have session_id |
328 | 330 | assert calls[1][1]["options"].params["session_id"] == "new_session" |
| 331 | + |
| 332 | + |
| 333 | +def test_code_interpreter_run_with_files(mocker): |
| 334 | + |
| 335 | + mock_requestor = mocker.MagicMock() |
| 336 | + response_data = { |
| 337 | + "data": { |
| 338 | + "session_id": "test_session_files", |
| 339 | + "status": "success", |
| 340 | + "outputs": [{"type": "stdout", "data": "File content read"}], |
| 341 | + } |
| 342 | + } |
| 343 | + mock_headers = { |
| 344 | + "cf-ray": "test-ray-id-files", |
| 345 | + "x-ratelimit-remaining": "98", |
| 346 | + "x-hostname": "test-host", |
| 347 | + "x-total-time": "42.0", |
| 348 | + } |
| 349 | + mock_response = TogetherResponse(data=response_data, headers=mock_headers) |
| 350 | + mock_requestor.request.return_value = (mock_response, None, None) |
| 351 | + mocker.patch( |
| 352 | + "together.abstract.api_requestor.APIRequestor", return_value=mock_requestor |
| 353 | + ) |
| 354 | + |
| 355 | + # Create code interpreter instance |
| 356 | + client = mocker.MagicMock() |
| 357 | + interpreter = CodeInterpreter(client) |
| 358 | + |
| 359 | + # Define files |
| 360 | + files_to_upload = [ |
| 361 | + {"name": "test.txt", "encoding": "string", "content": "Hello from file!"}, |
| 362 | + {"name": "image.png", "encoding": "base64", "content": "aW1hZ2UgZGF0YQ=="}, |
| 363 | + ] |
| 364 | + |
| 365 | + # Test run method with files (passing list of dicts) |
| 366 | + response = interpreter.run( |
| 367 | + code='with open("test.txt") as f: print(f.read())', |
| 368 | + language="python", |
| 369 | + files=files_to_upload, # Pass the list of dictionaries directly |
| 370 | + ) |
| 371 | + |
| 372 | + # Verify the response |
| 373 | + assert isinstance(response, ExecuteResponse) |
| 374 | + assert response.data.session_id == "test_session_files" |
| 375 | + assert response.data.status == "success" |
| 376 | + assert len(response.data.outputs) == 1 |
| 377 | + assert response.data.outputs[0].type == "stdout" |
| 378 | + |
| 379 | + # Verify API request includes files (expected_files_payload remains the same) |
| 380 | + mock_requestor.request.assert_called_once_with( |
| 381 | + options=mocker.ANY, |
| 382 | + stream=False, |
| 383 | + ) |
| 384 | + request_options = mock_requestor.request.call_args[1]["options"] |
| 385 | + assert request_options.method == "POST" |
| 386 | + assert request_options.url == "/tci/execute" |
| 387 | + expected_files_payload = [ |
| 388 | + {"name": "test.txt", "encoding": "string", "content": "Hello from file!"}, |
| 389 | + {"name": "image.png", "encoding": "base64", "content": "aW1hZ2UgZGF0YQ=="}, |
| 390 | + ] |
| 391 | + assert request_options.params == { |
| 392 | + "code": 'with open("test.txt") as f: print(f.read())', |
| 393 | + "language": "python", |
| 394 | + "files": expected_files_payload, |
| 395 | + } |
| 396 | + |
| 397 | + |
| 398 | +def test_code_interpreter_run_with_invalid_file_dict_structure(mocker): |
| 399 | + """Test that run raises ValueError for missing keys in file dict.""" |
| 400 | + client = mocker.MagicMock() |
| 401 | + interpreter = CodeInterpreter(client) |
| 402 | + |
| 403 | + invalid_files = [ |
| 404 | + {"name": "test.txt", "content": "Missing encoding"} # Missing 'encoding' |
| 405 | + ] |
| 406 | + |
| 407 | + with pytest.raises(ValueError, match="Invalid file input format"): |
| 408 | + interpreter.run( |
| 409 | + code="print('test')", |
| 410 | + language="python", |
| 411 | + files=invalid_files, |
| 412 | + ) |
| 413 | + |
| 414 | + |
| 415 | +def test_code_interpreter_run_with_invalid_file_dict_encoding(mocker): |
| 416 | + """Test that run raises ValueError for invalid encoding value.""" |
| 417 | + client = mocker.MagicMock() |
| 418 | + interpreter = CodeInterpreter(client) |
| 419 | + |
| 420 | + invalid_files = [ |
| 421 | + { |
| 422 | + "name": "test.txt", |
| 423 | + "encoding": "utf-8", |
| 424 | + "content": "Invalid encoding", |
| 425 | + } # Invalid 'encoding' value |
| 426 | + ] |
| 427 | + |
| 428 | + with pytest.raises(ValueError, match="Invalid file input format"): |
| 429 | + interpreter.run( |
| 430 | + code="print('test')", |
| 431 | + language="python", |
| 432 | + files=invalid_files, |
| 433 | + ) |
| 434 | + |
| 435 | + |
| 436 | +def test_code_interpreter_run_with_invalid_file_list_item(mocker): |
| 437 | + """Test that run raises ValueError for non-dict item in files list.""" |
| 438 | + client = mocker.MagicMock() |
| 439 | + interpreter = CodeInterpreter(client) |
| 440 | + |
| 441 | + invalid_files = [ |
| 442 | + {"name": "good.txt", "encoding": "string", "content": "Good"}, |
| 443 | + "not a dictionary", # Invalid item type |
| 444 | + ] |
| 445 | + |
| 446 | + with pytest.raises( |
| 447 | + ValueError, |
| 448 | + match="Invalid file input: Each item in 'files' must be a dictionary", |
| 449 | + ): |
| 450 | + interpreter.run( |
| 451 | + code="print('test')", |
| 452 | + language="python", |
| 453 | + files=invalid_files, |
| 454 | + ) |
0 commit comments