Skip to content

Commit

Permalink
Increase test coverage for output to file (#1049)
Browse files Browse the repository at this point in the history
  • Loading branch information
thiagomoretto authored Aug 10, 2024
1 parent 6ca051e commit b074138
Showing 1 changed file with 14 additions and 8 deletions.
22 changes: 14 additions & 8 deletions tests/task_test.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Test Agent creation and execution basic functionality."""

import os
import hashlib
import json
from unittest.mock import MagicMock, patch
Expand Down Expand Up @@ -277,6 +278,7 @@ class ScoreOutput(BaseModel):
description="Give me an integer score between 1-5 for the following title: 'The impact of AI in the future of work'",
expected_output="The score of the title.",
output_json=ScoreOutput,
output_file="score.json",
agent=scorer,
)

Expand Down Expand Up @@ -519,11 +521,13 @@ class ScoreOutput(BaseModel):
)

crew = Crew(agents=[scorer], tasks=[task])
crew.kickoff()

with patch.object(Task, "_save_file") as save_file:
save_file.return_value = None
crew.kickoff()
save_file.assert_called_once_with({"score": 4})
output_file_exists = os.path.exists("score.json")
assert output_file_exists
assert {"score": 4} == json.loads(open("score.json").read())
if output_file_exists:
os.remove("score.json")


@pytest.mark.vcr(filter_headers=["authorization"])
Expand All @@ -547,11 +551,13 @@ class ScoreOutput(BaseModel):
)

crew = Crew(agents=[scorer], tasks=[task])
crew.kickoff()

with patch.object(Task, "_save_file") as save_file:
save_file.return_value = None
crew.kickoff()
save_file.assert_called_once_with('{"score":4}')
output_file_exists = os.path.exists("score.json")
assert output_file_exists
assert {"score": 4} == json.loads(open("score.json").read())
if output_file_exists:
os.remove("score.json")


@pytest.mark.vcr(filter_headers=["authorization"])
Expand Down

0 comments on commit b074138

Please sign in to comment.