Spaces:
Running
Running
File size: 2,999 Bytes
9a255bf |
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 |
import unittest
from unittest.mock import patch, MagicMock, AsyncMock
import pandas as pd
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tinyagent import TinyCodeAgent
from matching_agent import create_matching_agent, run_matching
class TestMatchingAgent(unittest.TestCase):
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_key"})
def test_01_create_matching_agent(self):
"""
Test the creation and configuration of the matching agent.
"""
agent = create_matching_agent()
self.assertIsInstance(agent, TinyCodeAgent)
self.assertEqual(agent.model, "gpt-4.1-mini")
self.assertIn("You are a brilliant hackathon team-matching AI.", agent.system_prompt)
self.assertIn("participants_df", agent.system_prompt)
self.assertIn("organizer_criteria", agent.system_prompt)
# Check code execution environment settings
self.assertIn("pandas", agent.pip_packages)
self.assertIn("scikit-learn", agent.pip_packages)
self.assertIn("pandas", agent.authorized_imports)
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_key"})
def test_02_run_matching(self):
"""
Test the run_matching function to ensure it configures the agent
and calls the run method correctly.
"""
# Create a mock agent to isolate the test from the actual TinyCodeAgent implementation
mock_agent = MagicMock(spec=TinyCodeAgent)
mock_agent.run = AsyncMock(return_value="## Mocked Team Report")
# Sample data
participants_df = pd.DataFrame({
"name": ["Alice", "Bob"],
"skills": ["Frontend", "Backend"]
})
organizer_criteria = "Create teams of 2."
# Define an async test function to run the coroutine
async def do_run_matching():
result = await run_matching(mock_agent, participants_df, organizer_criteria)
# --- Assertions ---
# 1. Check that user_variables were set correctly on the agent
self.assertIn("participants_df", mock_agent.user_variables)
self.assertIn("organizer_criteria", mock_agent.user_variables)
pd.testing.assert_frame_equal(mock_agent.user_variables["participants_df"], participants_df)
self.assertEqual(mock_agent.user_variables["organizer_criteria"], organizer_criteria)
# 2. Check that agent.run was called correctly
mock_agent.run.assert_called_once()
# The first argument to run should be the user prompt
self.assertEqual(mock_agent.run.call_args[0][0], "Form the teams based on the provided data and criteria.")
# 3. Check the result
self.assertEqual(result, "## Mocked Team Report")
# Run the async test function
import asyncio
asyncio.run(do_run_matching())
if __name__ == '__main__':
unittest.main() |