File size: 2,224 Bytes
3b9a6b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import pytest
from datetime import datetime, timedelta
from src.utils.load_secrets import load_secrets
from dateutil.rrule import rrulestr
from icalendar import Calendar, vDDDTypes

# Load environment variables for agent (if needed)
load_secrets("tests/secrets/creds.py")

import factory.data_provider as data_provider
from src.utils.extract_calendar import extract_ical_entries


@pytest.mark.asyncio
async def test_factory_demo_agent():
    # Use a simple string as the project description
    test_input = "Test project for schedule generation."

    # Generate schedule data using generate_agent_data
    schedule = await data_provider.generate_agent_data(test_input)

    # Assert basic schedule properties
    assert len(schedule.employees) > 0
    assert schedule.schedule_info.total_slots > 0
    assert len(schedule.tasks) > 0

    # Verify employee skills
    for employee in schedule.employees:
        assert len(employee.skills) > 0
        # Check that each employee has at least one required skill
        assert any(
            skill in data_provider.SKILL_SET.required_skills
            for skill in employee.skills
        )

    # Verify task properties
    for task in schedule.tasks:
        assert task.duration_slots > 0
        assert task.required_skill
        assert hasattr(task, "project_id")

    # Print schedule details for debugging
    print("Employee names:", [e.name for e in schedule.employees])
    print("Tasks count:", len(schedule.tasks))
    print("Total slots:", schedule.schedule_info.total_slots)


@pytest.mark.asyncio
async def test_factory_mcp():
    # Load the real calendar.ics file
    with open("tests/data/calendar.ics", "rb") as f:
        file_bytes = f.read()
    entries, err = extract_ical_entries(file_bytes)
    assert err is None
    assert entries is not None
    assert len(entries) > 0

    print("\nEntries:")
    print(entries)

    # Use a made-up user message
    user_message = "Create a new AWS VPC."

    # Call generate_mcp_data directly
    df = await data_provider.generate_mcp_data(entries, user_message)

    # Assert the DataFrame is not empty
    assert df is not None
    assert not df.empty

    # Print the DataFrame for debug
    print(df)