test: add file utilities tests

- Add tests for file processing constraints and validators
- Add tests for FileProcessor and FileResolver
- Add tests for resolved file types
- Add tests for file store operations
- Add unit tests for multimodal LLM support
This commit is contained in:
Greyson LaLonde
2026-01-21 20:12:57 -05:00
parent b035aa8947
commit 1fe020fa6f
7 changed files with 1608 additions and 0 deletions

View File

@@ -0,0 +1,474 @@
"""Unit tests for LLM multimodal functionality across all providers."""
import base64
import os
from unittest.mock import patch
import pytest
from crewai.llm import LLM
from crewai.utilities.files import ImageFile, PDFFile, TextFile
# Check for optional provider dependencies
try:
from crewai.llms.providers.anthropic.completion import AnthropicCompletion
HAS_ANTHROPIC = True
except ImportError:
HAS_ANTHROPIC = False
try:
from crewai.llms.providers.azure.completion import AzureCompletion
HAS_AZURE = True
except ImportError:
HAS_AZURE = False
try:
from crewai.llms.providers.bedrock.completion import BedrockCompletion
HAS_BEDROCK = True
except ImportError:
HAS_BEDROCK = False
# Minimal valid PNG for testing
MINIMAL_PNG = (
b"\x89PNG\r\n\x1a\n"
b"\x00\x00\x00\rIHDR"
b"\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00"
b"\x90wS\xde"
b"\x00\x00\x00\x00IEND\xaeB`\x82"
)
MINIMAL_PDF = b"%PDF-1.4 test content"
@pytest.fixture(autouse=True)
def mock_api_keys():
"""Mock API keys for all providers."""
env_vars = {
"ANTHROPIC_API_KEY": "test-key",
"OPENAI_API_KEY": "test-key",
"GOOGLE_API_KEY": "test-key",
"AZURE_API_KEY": "test-key",
"AWS_ACCESS_KEY_ID": "test-key",
"AWS_SECRET_ACCESS_KEY": "test-key",
}
with patch.dict(os.environ, env_vars):
yield
class TestLiteLLMMultimodal:
"""Tests for LLM class (litellm wrapper) multimodal functionality.
These tests use `is_litellm=True` to ensure the litellm wrapper is used
instead of native providers.
"""
def test_supports_multimodal_gpt4o(self) -> None:
"""Test GPT-4o model supports multimodal."""
llm = LLM(model="gpt-4o", is_litellm=True)
assert llm.supports_multimodal() is True
def test_supports_multimodal_gpt4_turbo(self) -> None:
"""Test GPT-4 Turbo model supports multimodal."""
llm = LLM(model="gpt-4-turbo", is_litellm=True)
assert llm.supports_multimodal() is True
def test_supports_multimodal_claude3(self) -> None:
"""Test Claude 3 model supports multimodal via litellm."""
# Use litellm/ prefix to avoid native provider import
llm = LLM(model="litellm/claude-3-sonnet-20240229")
assert llm.supports_multimodal() is True
def test_supports_multimodal_gemini(self) -> None:
"""Test Gemini model supports multimodal."""
llm = LLM(model="gemini/gemini-pro", is_litellm=True)
assert llm.supports_multimodal() is True
def test_supports_multimodal_gpt35_does_not(self) -> None:
"""Test GPT-3.5 model does not support multimodal."""
llm = LLM(model="gpt-3.5-turbo", is_litellm=True)
assert llm.supports_multimodal() is False
def test_supported_content_types_openai(self) -> None:
"""Test OpenAI models support images only."""
llm = LLM(model="gpt-4o", is_litellm=True)
types = llm.supported_multimodal_content_types()
assert "image/" in types
assert "application/pdf" not in types
def test_supported_content_types_claude(self) -> None:
"""Test Claude models support images and PDFs via litellm."""
# Use litellm/ prefix to avoid native provider import
llm = LLM(model="litellm/claude-3-sonnet-20240229")
types = llm.supported_multimodal_content_types()
assert "image/" in types
assert "application/pdf" in types
def test_supported_content_types_gemini(self) -> None:
"""Test Gemini models support wide range of content."""
llm = LLM(model="gemini/gemini-pro", is_litellm=True)
types = llm.supported_multimodal_content_types()
assert "image/" in types
assert "audio/" in types
assert "video/" in types
assert "application/pdf" in types
assert "text/" in types
def test_supported_content_types_non_multimodal(self) -> None:
"""Test non-multimodal models return empty list."""
llm = LLM(model="gpt-3.5-turbo", is_litellm=True)
assert llm.supported_multimodal_content_types() == []
def test_format_multimodal_content_image(self) -> None:
"""Test formatting image content."""
llm = LLM(model="gpt-4o", is_litellm=True)
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert result[0]["type"] == "image_url"
assert "data:image/png;base64," in result[0]["image_url"]["url"]
def test_format_multimodal_content_non_multimodal(self) -> None:
"""Test non-multimodal model returns empty list."""
llm = LLM(model="gpt-3.5-turbo", is_litellm=True)
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert result == []
def test_format_multimodal_content_unsupported_type(self) -> None:
"""Test unsupported content type is skipped."""
llm = LLM(model="gpt-4o", is_litellm=True) # OpenAI doesn't support PDF
files = {"doc": PDFFile(source=MINIMAL_PDF)}
result = llm.format_multimodal_content(files)
assert result == []
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic SDK not installed")
class TestAnthropicMultimodal:
"""Tests for Anthropic provider multimodal functionality."""
def test_supports_multimodal_claude3(self) -> None:
"""Test Claude 3 supports multimodal."""
llm = LLM(model="anthropic/claude-3-sonnet-20240229")
assert llm.supports_multimodal() is True
def test_supports_multimodal_claude4(self) -> None:
"""Test Claude 4 supports multimodal."""
llm = LLM(model="anthropic/claude-4-opus")
assert llm.supports_multimodal() is True
def test_supported_content_types(self) -> None:
"""Test Anthropic supports images and PDFs."""
llm = LLM(model="anthropic/claude-3-sonnet-20240229")
types = llm.supported_multimodal_content_types()
assert "image/" in types
assert "application/pdf" in types
def test_format_multimodal_content_image(self) -> None:
"""Test Anthropic image format uses source-based structure."""
llm = LLM(model="anthropic/claude-3-sonnet-20240229")
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert result[0]["type"] == "image"
assert result[0]["source"]["type"] == "base64"
assert result[0]["source"]["media_type"] == "image/png"
assert "data" in result[0]["source"]
def test_format_multimodal_content_pdf(self) -> None:
"""Test Anthropic PDF format uses document structure."""
llm = LLM(model="anthropic/claude-3-sonnet-20240229")
files = {"doc": PDFFile(source=MINIMAL_PDF)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert result[0]["type"] == "document"
assert result[0]["source"]["type"] == "base64"
assert result[0]["source"]["media_type"] == "application/pdf"
class TestOpenAIMultimodal:
"""Tests for OpenAI provider multimodal functionality."""
def test_supports_multimodal_gpt4o(self) -> None:
"""Test GPT-4o supports multimodal."""
llm = LLM(model="openai/gpt-4o")
assert llm.supports_multimodal() is True
def test_supports_multimodal_gpt4_vision(self) -> None:
"""Test GPT-4 Vision supports multimodal."""
llm = LLM(model="openai/gpt-4-vision-preview")
assert llm.supports_multimodal() is True
def test_supports_multimodal_o1(self) -> None:
"""Test O1 model supports multimodal."""
llm = LLM(model="openai/o1-preview")
assert llm.supports_multimodal() is True
def test_does_not_support_gpt35(self) -> None:
"""Test GPT-3.5 does not support multimodal."""
llm = LLM(model="openai/gpt-3.5-turbo")
assert llm.supports_multimodal() is False
def test_supported_content_types(self) -> None:
"""Test OpenAI supports only images."""
llm = LLM(model="openai/gpt-4o")
types = llm.supported_multimodal_content_types()
assert types == ["image/"]
def test_format_multimodal_content_image(self) -> None:
"""Test OpenAI uses image_url format."""
llm = LLM(model="openai/gpt-4o")
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert result[0]["type"] == "image_url"
url = result[0]["image_url"]["url"]
assert url.startswith("data:image/png;base64,")
# Verify base64 content
b64_data = url.split(",")[1]
assert base64.b64decode(b64_data) == MINIMAL_PNG
class TestGeminiMultimodal:
"""Tests for Gemini provider multimodal functionality."""
def test_supports_multimodal_always_true(self) -> None:
"""Test Gemini always supports multimodal."""
llm = LLM(model="gemini/gemini-pro")
assert llm.supports_multimodal() is True
def test_supported_content_types(self) -> None:
"""Test Gemini supports wide range of types."""
llm = LLM(model="gemini/gemini-pro")
types = llm.supported_multimodal_content_types()
assert "image/" in types
assert "audio/" in types
assert "video/" in types
assert "application/pdf" in types
assert "text/" in types
def test_format_multimodal_content_image(self) -> None:
"""Test Gemini uses inlineData format."""
llm = LLM(model="gemini/gemini-pro")
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert "inlineData" in result[0]
assert result[0]["inlineData"]["mimeType"] == "image/png"
assert "data" in result[0]["inlineData"]
def test_format_text_content(self) -> None:
"""Test Gemini text format uses simple text key."""
llm = LLM(model="gemini/gemini-pro")
result = llm.format_text_content("Hello world")
assert result == {"text": "Hello world"}
@pytest.mark.skipif(not HAS_AZURE, reason="Azure AI Inference SDK not installed")
class TestAzureMultimodal:
"""Tests for Azure OpenAI provider multimodal functionality."""
@pytest.fixture(autouse=True)
def mock_azure_env(self):
"""Mock Azure-specific environment variables."""
env_vars = {
"AZURE_API_KEY": "test-key",
"AZURE_API_BASE": "https://test.openai.azure.com",
"AZURE_API_VERSION": "2024-02-01",
}
with patch.dict(os.environ, env_vars):
yield
def test_supports_multimodal_gpt4o(self) -> None:
"""Test Azure GPT-4o supports multimodal."""
llm = LLM(model="azure/gpt-4o")
assert llm.supports_multimodal() is True
def test_supports_multimodal_gpt4_turbo(self) -> None:
"""Test Azure GPT-4 Turbo supports multimodal."""
llm = LLM(model="azure/gpt-4-turbo")
assert llm.supports_multimodal() is True
def test_does_not_support_gpt35(self) -> None:
"""Test Azure GPT-3.5 does not support multimodal."""
llm = LLM(model="azure/gpt-35-turbo")
assert llm.supports_multimodal() is False
def test_supported_content_types(self) -> None:
"""Test Azure supports only images."""
llm = LLM(model="azure/gpt-4o")
types = llm.supported_multimodal_content_types()
assert types == ["image/"]
def test_format_multimodal_content_image(self) -> None:
"""Test Azure uses same format as OpenAI."""
llm = LLM(model="azure/gpt-4o")
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert result[0]["type"] == "image_url"
assert "data:image/png;base64," in result[0]["image_url"]["url"]
@pytest.mark.skipif(not HAS_BEDROCK, reason="AWS Bedrock SDK not installed")
class TestBedrockMultimodal:
"""Tests for AWS Bedrock provider multimodal functionality."""
@pytest.fixture(autouse=True)
def mock_bedrock_env(self):
"""Mock AWS-specific environment variables."""
env_vars = {
"AWS_ACCESS_KEY_ID": "test-key",
"AWS_SECRET_ACCESS_KEY": "test-secret",
"AWS_DEFAULT_REGION": "us-east-1",
}
with patch.dict(os.environ, env_vars):
yield
def test_supports_multimodal_claude3(self) -> None:
"""Test Bedrock Claude 3 supports multimodal."""
llm = LLM(model="bedrock/anthropic.claude-3-sonnet")
assert llm.supports_multimodal() is True
def test_does_not_support_claude2(self) -> None:
"""Test Bedrock Claude 2 does not support multimodal."""
llm = LLM(model="bedrock/anthropic.claude-v2")
assert llm.supports_multimodal() is False
def test_supported_content_types(self) -> None:
"""Test Bedrock supports images and PDFs."""
llm = LLM(model="bedrock/anthropic.claude-3-sonnet")
types = llm.supported_multimodal_content_types()
assert "image/" in types
assert "application/pdf" in types
def test_format_multimodal_content_image(self) -> None:
"""Test Bedrock uses Converse API image format."""
llm = LLM(model="bedrock/anthropic.claude-3-sonnet")
files = {"chart": ImageFile(source=MINIMAL_PNG)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert "image" in result[0]
assert result[0]["image"]["format"] == "png"
assert "source" in result[0]["image"]
assert "bytes" in result[0]["image"]["source"]
def test_format_multimodal_content_pdf(self) -> None:
"""Test Bedrock uses Converse API document format."""
llm = LLM(model="bedrock/anthropic.claude-3-sonnet")
files = {"doc": PDFFile(source=MINIMAL_PDF)}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert "document" in result[0]
assert result[0]["document"]["format"] == "pdf"
assert "source" in result[0]["document"]
class TestBaseLLMMultimodal:
"""Tests for BaseLLM default multimodal behavior."""
def test_base_supports_multimodal_false(self) -> None:
"""Test base implementation returns False."""
from crewai.llms.base_llm import BaseLLM
class TestLLM(BaseLLM):
def call(self, messages, tools=None, callbacks=None):
return "test"
llm = TestLLM(model="test")
assert llm.supports_multimodal() is False
def test_base_supported_content_types_empty(self) -> None:
"""Test base implementation returns empty list."""
from crewai.llms.base_llm import BaseLLM
class TestLLM(BaseLLM):
def call(self, messages, tools=None, callbacks=None):
return "test"
llm = TestLLM(model="test")
assert llm.supported_multimodal_content_types() == []
def test_base_format_multimodal_content_empty(self) -> None:
"""Test base implementation returns empty list."""
from crewai.llms.base_llm import BaseLLM
class TestLLM(BaseLLM):
def call(self, messages, tools=None, callbacks=None):
return "test"
llm = TestLLM(model="test")
files = {"chart": ImageFile(source=MINIMAL_PNG)}
assert llm.format_multimodal_content(files) == []
def test_base_format_text_content(self) -> None:
"""Test base text formatting uses OpenAI/Anthropic style."""
from crewai.llms.base_llm import BaseLLM
class TestLLM(BaseLLM):
def call(self, messages, tools=None, callbacks=None):
return "test"
llm = TestLLM(model="test")
result = llm.format_text_content("Hello")
assert result == {"type": "text", "text": "Hello"}
class TestMultipleFilesFormatting:
"""Tests for formatting multiple files at once."""
def test_format_multiple_images(self) -> None:
"""Test formatting multiple images."""
llm = LLM(model="gpt-4o")
files = {
"chart1": ImageFile(source=MINIMAL_PNG),
"chart2": ImageFile(source=MINIMAL_PNG),
}
result = llm.format_multimodal_content(files)
assert len(result) == 2
def test_format_mixed_supported_and_unsupported(self) -> None:
"""Test only supported types are formatted."""
llm = LLM(model="gpt-4o") # OpenAI - images only
files = {
"chart": ImageFile(source=MINIMAL_PNG),
"doc": PDFFile(source=MINIMAL_PDF), # Not supported
"text": TextFile(source=b"hello"), # Not supported
}
result = llm.format_multimodal_content(files)
assert len(result) == 1
assert result[0]["type"] == "image_url"
def test_format_empty_files_dict(self) -> None:
"""Test empty files dict returns empty list."""
llm = LLM(model="gpt-4o")
result = llm.format_multimodal_content({})
assert result == []

View File

@@ -0,0 +1,226 @@
"""Tests for provider constraints."""
import pytest
from crewai.utilities.files.processing.constraints import (
ANTHROPIC_CONSTRAINTS,
BEDROCK_CONSTRAINTS,
GEMINI_CONSTRAINTS,
OPENAI_CONSTRAINTS,
AudioConstraints,
ImageConstraints,
PDFConstraints,
ProviderConstraints,
VideoConstraints,
get_constraints_for_provider,
)
class TestImageConstraints:
"""Tests for ImageConstraints dataclass."""
def test_image_constraints_creation(self):
"""Test creating image constraints with all fields."""
constraints = ImageConstraints(
max_size_bytes=5 * 1024 * 1024,
max_width=8000,
max_height=8000,
max_images_per_request=10,
)
assert constraints.max_size_bytes == 5 * 1024 * 1024
assert constraints.max_width == 8000
assert constraints.max_height == 8000
assert constraints.max_images_per_request == 10
def test_image_constraints_defaults(self):
"""Test image constraints with default values."""
constraints = ImageConstraints(max_size_bytes=1000)
assert constraints.max_size_bytes == 1000
assert constraints.max_width is None
assert constraints.max_height is None
assert constraints.max_images_per_request is None
assert "image/png" in constraints.supported_formats
def test_image_constraints_frozen(self):
"""Test that image constraints are immutable."""
constraints = ImageConstraints(max_size_bytes=1000)
with pytest.raises(Exception):
constraints.max_size_bytes = 2000
class TestPDFConstraints:
"""Tests for PDFConstraints dataclass."""
def test_pdf_constraints_creation(self):
"""Test creating PDF constraints."""
constraints = PDFConstraints(
max_size_bytes=30 * 1024 * 1024,
max_pages=100,
)
assert constraints.max_size_bytes == 30 * 1024 * 1024
assert constraints.max_pages == 100
def test_pdf_constraints_defaults(self):
"""Test PDF constraints with default values."""
constraints = PDFConstraints(max_size_bytes=1000)
assert constraints.max_size_bytes == 1000
assert constraints.max_pages is None
class TestAudioConstraints:
"""Tests for AudioConstraints dataclass."""
def test_audio_constraints_creation(self):
"""Test creating audio constraints."""
constraints = AudioConstraints(
max_size_bytes=100 * 1024 * 1024,
max_duration_seconds=3600,
)
assert constraints.max_size_bytes == 100 * 1024 * 1024
assert constraints.max_duration_seconds == 3600
assert "audio/mp3" in constraints.supported_formats
class TestVideoConstraints:
"""Tests for VideoConstraints dataclass."""
def test_video_constraints_creation(self):
"""Test creating video constraints."""
constraints = VideoConstraints(
max_size_bytes=2 * 1024 * 1024 * 1024,
max_duration_seconds=7200,
)
assert constraints.max_size_bytes == 2 * 1024 * 1024 * 1024
assert constraints.max_duration_seconds == 7200
assert "video/mp4" in constraints.supported_formats
class TestProviderConstraints:
"""Tests for ProviderConstraints dataclass."""
def test_provider_constraints_creation(self):
"""Test creating full provider constraints."""
constraints = ProviderConstraints(
name="test-provider",
image=ImageConstraints(max_size_bytes=5 * 1024 * 1024),
pdf=PDFConstraints(max_size_bytes=30 * 1024 * 1024),
supports_file_upload=True,
file_upload_threshold_bytes=10 * 1024 * 1024,
)
assert constraints.name == "test-provider"
assert constraints.image is not None
assert constraints.pdf is not None
assert constraints.supports_file_upload is True
def test_provider_constraints_defaults(self):
"""Test provider constraints with default values."""
constraints = ProviderConstraints(name="test")
assert constraints.name == "test"
assert constraints.image is None
assert constraints.pdf is None
assert constraints.audio is None
assert constraints.video is None
assert constraints.supports_file_upload is False
class TestPredefinedConstraints:
"""Tests for predefined provider constraints."""
def test_anthropic_constraints(self):
"""Test Anthropic constraints are properly defined."""
assert ANTHROPIC_CONSTRAINTS.name == "anthropic"
assert ANTHROPIC_CONSTRAINTS.image is not None
assert ANTHROPIC_CONSTRAINTS.image.max_size_bytes == 5 * 1024 * 1024
assert ANTHROPIC_CONSTRAINTS.image.max_width == 8000
assert ANTHROPIC_CONSTRAINTS.pdf is not None
assert ANTHROPIC_CONSTRAINTS.pdf.max_pages == 100
assert ANTHROPIC_CONSTRAINTS.supports_file_upload is True
def test_openai_constraints(self):
"""Test OpenAI constraints are properly defined."""
assert OPENAI_CONSTRAINTS.name == "openai"
assert OPENAI_CONSTRAINTS.image is not None
assert OPENAI_CONSTRAINTS.image.max_size_bytes == 20 * 1024 * 1024
assert OPENAI_CONSTRAINTS.pdf is None # OpenAI doesn't support PDFs
def test_gemini_constraints(self):
"""Test Gemini constraints are properly defined."""
assert GEMINI_CONSTRAINTS.name == "gemini"
assert GEMINI_CONSTRAINTS.image is not None
assert GEMINI_CONSTRAINTS.pdf is not None
assert GEMINI_CONSTRAINTS.audio is not None
assert GEMINI_CONSTRAINTS.video is not None
assert GEMINI_CONSTRAINTS.supports_file_upload is True
def test_bedrock_constraints(self):
"""Test Bedrock constraints are properly defined."""
assert BEDROCK_CONSTRAINTS.name == "bedrock"
assert BEDROCK_CONSTRAINTS.image is not None
assert BEDROCK_CONSTRAINTS.image.max_size_bytes == 4_608_000
assert BEDROCK_CONSTRAINTS.pdf is not None
assert BEDROCK_CONSTRAINTS.supports_file_upload is False
class TestGetConstraintsForProvider:
"""Tests for get_constraints_for_provider function."""
def test_get_by_exact_name(self):
"""Test getting constraints by exact provider name."""
result = get_constraints_for_provider("anthropic")
assert result == ANTHROPIC_CONSTRAINTS
result = get_constraints_for_provider("openai")
assert result == OPENAI_CONSTRAINTS
result = get_constraints_for_provider("gemini")
assert result == GEMINI_CONSTRAINTS
def test_get_by_alias(self):
"""Test getting constraints by alias name."""
result = get_constraints_for_provider("claude")
assert result == ANTHROPIC_CONSTRAINTS
result = get_constraints_for_provider("gpt")
assert result == OPENAI_CONSTRAINTS
result = get_constraints_for_provider("google")
assert result == GEMINI_CONSTRAINTS
def test_get_case_insensitive(self):
"""Test case-insensitive lookup."""
result = get_constraints_for_provider("ANTHROPIC")
assert result == ANTHROPIC_CONSTRAINTS
result = get_constraints_for_provider("OpenAI")
assert result == OPENAI_CONSTRAINTS
def test_get_with_provider_constraints_object(self):
"""Test passing ProviderConstraints object returns it unchanged."""
custom = ProviderConstraints(name="custom")
result = get_constraints_for_provider(custom)
assert result is custom
def test_get_unknown_provider(self):
"""Test unknown provider returns None."""
result = get_constraints_for_provider("unknown-provider")
assert result is None
def test_get_by_partial_match(self):
"""Test partial match in provider string."""
result = get_constraints_for_provider("claude-3-sonnet")
assert result == ANTHROPIC_CONSTRAINTS
result = get_constraints_for_provider("gpt-4o")
assert result == OPENAI_CONSTRAINTS
result = get_constraints_for_provider("gemini-pro")
assert result == GEMINI_CONSTRAINTS

View File

@@ -0,0 +1,220 @@
"""Tests for FileProcessor class."""
import pytest
from crewai.utilities.files import FileBytes, ImageFile, PDFFile, TextFile
from crewai.utilities.files.processing.constraints import (
ANTHROPIC_CONSTRAINTS,
ImageConstraints,
PDFConstraints,
ProviderConstraints,
)
from crewai.utilities.files.processing.enums import FileHandling
from crewai.utilities.files.processing.exceptions import (
FileTooLargeError,
FileValidationError,
)
from crewai.utilities.files.processing.processor import FileProcessor
# Minimal valid PNG: 8x8 pixel RGB image (valid for PIL)
MINIMAL_PNG = bytes([
0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08,
0x08, 0x02, 0x00, 0x00, 0x00, 0x4b, 0x6d, 0x29, 0xdc, 0x00, 0x00, 0x00,
0x12, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x63, 0xfc, 0xcf, 0x80, 0x1d,
0x30, 0xe1, 0x10, 0x1f, 0xa4, 0x12, 0x00, 0xcd, 0x41, 0x01, 0x0f, 0xe8,
0x41, 0xe2, 0x6f, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae,
0x42, 0x60, 0x82,
])
# Minimal valid PDF
MINIMAL_PDF = (
b"%PDF-1.4\n1 0 obj<</Type/Catalog/Pages 2 0 R>>endobj "
b"2 0 obj<</Type/Pages/Kids[3 0 R]/Count 1>>endobj "
b"3 0 obj<</Type/Page/MediaBox[0 0 612 792]/Parent 2 0 R>>endobj "
b"xref\n0 4\n0000000000 65535 f \n0000000009 00000 n \n"
b"0000000052 00000 n \n0000000101 00000 n \n"
b"trailer<</Size 4/Root 1 0 R>>\nstartxref\n178\n%%EOF"
)
class TestFileProcessorInit:
"""Tests for FileProcessor initialization."""
def test_init_with_constraints(self):
"""Test initialization with ProviderConstraints."""
processor = FileProcessor(constraints=ANTHROPIC_CONSTRAINTS)
assert processor.constraints == ANTHROPIC_CONSTRAINTS
def test_init_with_provider_string(self):
"""Test initialization with provider name string."""
processor = FileProcessor(constraints="anthropic")
assert processor.constraints == ANTHROPIC_CONSTRAINTS
def test_init_with_unknown_provider(self):
"""Test initialization with unknown provider sets constraints to None."""
processor = FileProcessor(constraints="unknown")
assert processor.constraints is None
def test_init_with_none_constraints(self):
"""Test initialization with None constraints."""
processor = FileProcessor(constraints=None)
assert processor.constraints is None
class TestFileProcessorValidate:
"""Tests for FileProcessor.validate method."""
def test_validate_valid_file(self):
"""Test validating a valid file returns no errors."""
processor = FileProcessor(constraints=ANTHROPIC_CONSTRAINTS)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
errors = processor.validate(file)
assert len(errors) == 0
def test_validate_without_constraints(self):
"""Test validating without constraints returns empty list."""
processor = FileProcessor(constraints=None)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
errors = processor.validate(file)
assert len(errors) == 0
def test_validate_strict_raises_on_error(self):
"""Test STRICT mode raises on validation error."""
constraints = ProviderConstraints(
name="test",
image=ImageConstraints(max_size_bytes=10),
)
processor = FileProcessor(constraints=constraints)
# Set mode to strict on the file
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"), mode="strict")
with pytest.raises(FileTooLargeError):
processor.validate(file)
class TestFileProcessorProcess:
"""Tests for FileProcessor.process method."""
def test_process_valid_file(self):
"""Test processing a valid file returns it unchanged."""
processor = FileProcessor(constraints=ANTHROPIC_CONSTRAINTS)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
result = processor.process(file)
assert result == file
def test_process_without_constraints(self):
"""Test processing without constraints returns file unchanged."""
processor = FileProcessor(constraints=None)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
result = processor.process(file)
assert result == file
def test_process_strict_raises_on_error(self):
"""Test STRICT mode raises on processing error."""
constraints = ProviderConstraints(
name="test",
image=ImageConstraints(max_size_bytes=10),
)
processor = FileProcessor(constraints=constraints)
# Set mode to strict on the file
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"), mode="strict")
with pytest.raises(FileTooLargeError):
processor.process(file)
def test_process_warn_returns_file(self):
"""Test WARN mode returns file with warning."""
constraints = ProviderConstraints(
name="test",
image=ImageConstraints(max_size_bytes=10),
)
processor = FileProcessor(constraints=constraints)
# Set mode to warn on the file
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"), mode="warn")
result = processor.process(file)
assert result == file
class TestFileProcessorProcessFiles:
"""Tests for FileProcessor.process_files method."""
def test_process_files_multiple(self):
"""Test processing multiple files."""
processor = FileProcessor(constraints=ANTHROPIC_CONSTRAINTS)
files = {
"image1": ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test1.png")),
"image2": ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test2.png")),
}
result = processor.process_files(files)
assert len(result) == 2
assert "image1" in result
assert "image2" in result
def test_process_files_empty(self):
"""Test processing empty files dict."""
processor = FileProcessor(constraints=ANTHROPIC_CONSTRAINTS)
result = processor.process_files({})
assert result == {}
class TestFileHandlingEnum:
"""Tests for FileHandling enum."""
def test_enum_values(self):
"""Test all enum values are accessible."""
assert FileHandling.STRICT.value == "strict"
assert FileHandling.AUTO.value == "auto"
assert FileHandling.WARN.value == "warn"
assert FileHandling.CHUNK.value == "chunk"
class TestFileProcessorPerFileMode:
"""Tests for per-file mode handling."""
def test_file_default_mode_is_auto(self):
"""Test that files default to auto mode."""
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
assert file.mode == "auto"
def test_file_custom_mode(self):
"""Test setting custom mode on file."""
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"), mode="strict")
assert file.mode == "strict"
def test_processor_respects_file_mode(self):
"""Test processor uses each file's mode setting."""
constraints = ProviderConstraints(
name="test",
image=ImageConstraints(max_size_bytes=10),
)
processor = FileProcessor(constraints=constraints)
# File with strict mode should raise
strict_file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"), mode="strict")
with pytest.raises(FileTooLargeError):
processor.process(strict_file)
# File with warn mode should not raise
warn_file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"), mode="warn")
result = processor.process(warn_file)
assert result == warn_file

View File

@@ -0,0 +1,208 @@
"""Tests for file validators."""
import pytest
from crewai.utilities.files import FileBytes, ImageFile, PDFFile, TextFile
from crewai.utilities.files.processing.constraints import (
ANTHROPIC_CONSTRAINTS,
ImageConstraints,
PDFConstraints,
ProviderConstraints,
)
from crewai.utilities.files.processing.exceptions import (
FileTooLargeError,
FileValidationError,
UnsupportedFileTypeError,
)
from crewai.utilities.files.processing.validators import (
validate_file,
validate_image,
validate_pdf,
validate_text,
)
# Minimal valid PNG: 8x8 pixel RGB image (valid for PIL)
MINIMAL_PNG = bytes([
0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08,
0x08, 0x02, 0x00, 0x00, 0x00, 0x4b, 0x6d, 0x29, 0xdc, 0x00, 0x00, 0x00,
0x12, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x63, 0xfc, 0xcf, 0x80, 0x1d,
0x30, 0xe1, 0x10, 0x1f, 0xa4, 0x12, 0x00, 0xcd, 0x41, 0x01, 0x0f, 0xe8,
0x41, 0xe2, 0x6f, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae,
0x42, 0x60, 0x82,
])
# Minimal valid PDF
MINIMAL_PDF = (
b"%PDF-1.4\n1 0 obj<</Type/Catalog/Pages 2 0 R>>endobj "
b"2 0 obj<</Type/Pages/Kids[3 0 R]/Count 1>>endobj "
b"3 0 obj<</Type/Page/MediaBox[0 0 612 792]/Parent 2 0 R>>endobj "
b"xref\n0 4\n0000000000 65535 f \n0000000009 00000 n \n"
b"0000000052 00000 n \n0000000101 00000 n \n"
b"trailer<</Size 4/Root 1 0 R>>\nstartxref\n178\n%%EOF"
)
class TestValidateImage:
"""Tests for validate_image function."""
def test_validate_valid_image(self):
"""Test validating a valid image within constraints."""
constraints = ImageConstraints(
max_size_bytes=10 * 1024 * 1024,
supported_formats=("image/png",),
)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
errors = validate_image(file, constraints, raise_on_error=False)
assert len(errors) == 0
def test_validate_image_too_large(self):
"""Test validating an image that exceeds size limit."""
constraints = ImageConstraints(
max_size_bytes=10, # Very small limit
supported_formats=("image/png",),
)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
with pytest.raises(FileTooLargeError) as exc_info:
validate_image(file, constraints)
assert "exceeds" in str(exc_info.value)
assert exc_info.value.file_name == "test.png"
def test_validate_image_unsupported_format(self):
"""Test validating an image with unsupported format."""
constraints = ImageConstraints(
max_size_bytes=10 * 1024 * 1024,
supported_formats=("image/jpeg",), # Only JPEG
)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
with pytest.raises(UnsupportedFileTypeError) as exc_info:
validate_image(file, constraints)
assert "not supported" in str(exc_info.value)
def test_validate_image_no_raise(self):
"""Test validating with raise_on_error=False returns errors list."""
constraints = ImageConstraints(
max_size_bytes=10,
supported_formats=("image/jpeg",),
)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
errors = validate_image(file, constraints, raise_on_error=False)
assert len(errors) == 2 # Size error and format error
class TestValidatePDF:
"""Tests for validate_pdf function."""
def test_validate_valid_pdf(self):
"""Test validating a valid PDF within constraints."""
constraints = PDFConstraints(
max_size_bytes=10 * 1024 * 1024,
)
file = PDFFile(source=FileBytes(data=MINIMAL_PDF, filename="test.pdf"))
errors = validate_pdf(file, constraints, raise_on_error=False)
assert len(errors) == 0
def test_validate_pdf_too_large(self):
"""Test validating a PDF that exceeds size limit."""
constraints = PDFConstraints(
max_size_bytes=10, # Very small limit
)
file = PDFFile(source=FileBytes(data=MINIMAL_PDF, filename="test.pdf"))
with pytest.raises(FileTooLargeError) as exc_info:
validate_pdf(file, constraints)
assert "exceeds" in str(exc_info.value)
class TestValidateText:
"""Tests for validate_text function."""
def test_validate_valid_text(self):
"""Test validating a valid text file."""
constraints = ProviderConstraints(
name="test",
general_max_size_bytes=10 * 1024 * 1024,
)
file = TextFile(source=FileBytes(data=b"Hello, World!", filename="test.txt"))
errors = validate_text(file, constraints, raise_on_error=False)
assert len(errors) == 0
def test_validate_text_too_large(self):
"""Test validating text that exceeds size limit."""
constraints = ProviderConstraints(
name="test",
general_max_size_bytes=5,
)
file = TextFile(source=FileBytes(data=b"Hello, World!", filename="test.txt"))
with pytest.raises(FileTooLargeError):
validate_text(file, constraints)
def test_validate_text_no_limit(self):
"""Test validating text with no size limit."""
constraints = ProviderConstraints(name="test")
file = TextFile(source=FileBytes(data=b"Hello, World!", filename="test.txt"))
errors = validate_text(file, constraints, raise_on_error=False)
assert len(errors) == 0
class TestValidateFile:
"""Tests for validate_file function."""
def test_validate_file_dispatches_to_image(self):
"""Test validate_file dispatches to image validator."""
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
errors = validate_file(file, ANTHROPIC_CONSTRAINTS, raise_on_error=False)
assert len(errors) == 0
def test_validate_file_dispatches_to_pdf(self):
"""Test validate_file dispatches to PDF validator."""
file = PDFFile(source=FileBytes(data=MINIMAL_PDF, filename="test.pdf"))
errors = validate_file(file, ANTHROPIC_CONSTRAINTS, raise_on_error=False)
assert len(errors) == 0
def test_validate_file_unsupported_type(self):
"""Test validating a file type not supported by provider."""
constraints = ProviderConstraints(
name="test",
image=None, # No image support
)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
with pytest.raises(UnsupportedFileTypeError) as exc_info:
validate_file(file, constraints)
assert "does not support images" in str(exc_info.value)
def test_validate_file_pdf_not_supported(self):
"""Test validating PDF when provider doesn't support it."""
constraints = ProviderConstraints(
name="test",
pdf=None, # No PDF support
)
file = PDFFile(source=FileBytes(data=MINIMAL_PDF, filename="test.pdf"))
with pytest.raises(UnsupportedFileTypeError) as exc_info:
validate_file(file, constraints)
assert "does not support PDFs" in str(exc_info.value)

View File

@@ -0,0 +1,135 @@
"""Tests for resolved file types."""
from datetime import datetime, timezone
import pytest
from crewai.utilities.files.resolved import (
FileReference,
InlineBase64,
InlineBytes,
ResolvedFile,
UrlReference,
)
class TestInlineBase64:
"""Tests for InlineBase64 resolved type."""
def test_create_inline_base64(self):
"""Test creating InlineBase64 instance."""
resolved = InlineBase64(
content_type="image/png",
data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
)
assert resolved.content_type == "image/png"
assert len(resolved.data) > 0
def test_inline_base64_is_resolved_file(self):
"""Test InlineBase64 is a ResolvedFile."""
resolved = InlineBase64(content_type="image/png", data="abc123")
assert isinstance(resolved, ResolvedFile)
def test_inline_base64_frozen(self):
"""Test InlineBase64 is immutable."""
resolved = InlineBase64(content_type="image/png", data="abc123")
with pytest.raises(Exception):
resolved.data = "xyz789"
class TestInlineBytes:
"""Tests for InlineBytes resolved type."""
def test_create_inline_bytes(self):
"""Test creating InlineBytes instance."""
data = b"\x89PNG\r\n\x1a\n"
resolved = InlineBytes(
content_type="image/png",
data=data,
)
assert resolved.content_type == "image/png"
assert resolved.data == data
def test_inline_bytes_is_resolved_file(self):
"""Test InlineBytes is a ResolvedFile."""
resolved = InlineBytes(content_type="image/png", data=b"test")
assert isinstance(resolved, ResolvedFile)
class TestFileReference:
"""Tests for FileReference resolved type."""
def test_create_file_reference(self):
"""Test creating FileReference instance."""
resolved = FileReference(
content_type="image/png",
file_id="file-abc123",
provider="gemini",
)
assert resolved.content_type == "image/png"
assert resolved.file_id == "file-abc123"
assert resolved.provider == "gemini"
assert resolved.expires_at is None
assert resolved.file_uri is None
def test_file_reference_with_expiry(self):
"""Test FileReference with expiry time."""
expiry = datetime.now(timezone.utc)
resolved = FileReference(
content_type="application/pdf",
file_id="file-xyz789",
provider="gemini",
expires_at=expiry,
)
assert resolved.expires_at == expiry
def test_file_reference_with_uri(self):
"""Test FileReference with URI."""
resolved = FileReference(
content_type="video/mp4",
file_id="file-video123",
provider="gemini",
file_uri="https://generativelanguage.googleapis.com/v1/files/file-video123",
)
assert resolved.file_uri is not None
def test_file_reference_is_resolved_file(self):
"""Test FileReference is a ResolvedFile."""
resolved = FileReference(
content_type="image/png",
file_id="file-123",
provider="anthropic",
)
assert isinstance(resolved, ResolvedFile)
class TestUrlReference:
"""Tests for UrlReference resolved type."""
def test_create_url_reference(self):
"""Test creating UrlReference instance."""
resolved = UrlReference(
content_type="image/png",
url="https://storage.googleapis.com/bucket/image.png",
)
assert resolved.content_type == "image/png"
assert resolved.url == "https://storage.googleapis.com/bucket/image.png"
def test_url_reference_is_resolved_file(self):
"""Test UrlReference is a ResolvedFile."""
resolved = UrlReference(
content_type="image/jpeg",
url="https://example.com/photo.jpg",
)
assert isinstance(resolved, ResolvedFile)

View File

@@ -0,0 +1,174 @@
"""Tests for FileResolver."""
import pytest
from crewai.utilities.files import FileBytes, ImageFile
from crewai.utilities.files.resolved import InlineBase64, InlineBytes
from crewai.utilities.files.resolver import (
FileResolver,
FileResolverConfig,
create_resolver,
)
from crewai.utilities.files.upload_cache import UploadCache
# Minimal valid PNG
MINIMAL_PNG = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x08\x00\x00\x00\x08"
b"\x01\x00\x00\x00\x00\xf9Y\xab\xcd\x00\x00\x00\nIDATx\x9cc`\x00\x00"
b"\x00\x02\x00\x01\xe2!\xbc3\x00\x00\x00\x00IEND\xaeB`\x82"
)
class TestFileResolverConfig:
"""Tests for FileResolverConfig."""
def test_default_config(self):
"""Test default configuration values."""
config = FileResolverConfig()
assert config.prefer_upload is False
assert config.upload_threshold_bytes is None
assert config.use_bytes_for_bedrock is True
def test_custom_config(self):
"""Test custom configuration values."""
config = FileResolverConfig(
prefer_upload=True,
upload_threshold_bytes=1024 * 1024,
use_bytes_for_bedrock=False,
)
assert config.prefer_upload is True
assert config.upload_threshold_bytes == 1024 * 1024
assert config.use_bytes_for_bedrock is False
class TestFileResolver:
"""Tests for FileResolver class."""
def test_resolve_inline_base64(self):
"""Test resolving file as inline base64."""
resolver = FileResolver()
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
resolved = resolver.resolve(file, "openai")
assert isinstance(resolved, InlineBase64)
assert resolved.content_type == "image/png"
assert len(resolved.data) > 0
def test_resolve_inline_bytes_for_bedrock(self):
"""Test resolving file as inline bytes for Bedrock."""
config = FileResolverConfig(use_bytes_for_bedrock=True)
resolver = FileResolver(config=config)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
resolved = resolver.resolve(file, "bedrock")
assert isinstance(resolved, InlineBytes)
assert resolved.content_type == "image/png"
assert resolved.data == MINIMAL_PNG
def test_resolve_files_multiple(self):
"""Test resolving multiple files."""
resolver = FileResolver()
files = {
"image1": ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test1.png")),
"image2": ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test2.png")),
}
resolved = resolver.resolve_files(files, "openai")
assert len(resolved) == 2
assert "image1" in resolved
assert "image2" in resolved
assert all(isinstance(r, InlineBase64) for r in resolved.values())
def test_resolve_with_cache(self):
"""Test resolver uses cache."""
cache = UploadCache()
resolver = FileResolver(upload_cache=cache)
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
# First resolution
resolved1 = resolver.resolve(file, "openai")
# Second resolution (should use same base64 encoding)
resolved2 = resolver.resolve(file, "openai")
assert isinstance(resolved1, InlineBase64)
assert isinstance(resolved2, InlineBase64)
# Data should be identical
assert resolved1.data == resolved2.data
def test_clear_cache(self):
"""Test clearing resolver cache."""
cache = UploadCache()
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
# Add something to cache manually
cache.set(file=file, provider="gemini", file_id="test")
resolver = FileResolver(upload_cache=cache)
resolver.clear_cache()
assert len(cache) == 0
def test_get_cached_uploads(self):
"""Test getting cached uploads from resolver."""
cache = UploadCache()
file = ImageFile(source=FileBytes(data=MINIMAL_PNG, filename="test.png"))
cache.set(file=file, provider="gemini", file_id="test-1")
cache.set(file=file, provider="anthropic", file_id="test-2")
resolver = FileResolver(upload_cache=cache)
gemini_uploads = resolver.get_cached_uploads("gemini")
anthropic_uploads = resolver.get_cached_uploads("anthropic")
assert len(gemini_uploads) == 1
assert len(anthropic_uploads) == 1
def test_get_cached_uploads_empty(self):
"""Test getting cached uploads when no cache."""
resolver = FileResolver() # No cache
uploads = resolver.get_cached_uploads("gemini")
assert uploads == []
class TestCreateResolver:
"""Tests for create_resolver factory function."""
def test_create_default_resolver(self):
"""Test creating resolver with default settings."""
resolver = create_resolver()
assert resolver.config.prefer_upload is False
assert resolver.upload_cache is not None
def test_create_resolver_with_options(self):
"""Test creating resolver with custom options."""
resolver = create_resolver(
prefer_upload=True,
upload_threshold_bytes=5 * 1024 * 1024,
enable_cache=False,
)
assert resolver.config.prefer_upload is True
assert resolver.config.upload_threshold_bytes == 5 * 1024 * 1024
assert resolver.upload_cache is None
def test_create_resolver_cache_enabled(self):
"""Test resolver has cache when enabled."""
resolver = create_resolver(enable_cache=True)
assert resolver.upload_cache is not None
def test_create_resolver_cache_disabled(self):
"""Test resolver has no cache when disabled."""
resolver = create_resolver(enable_cache=False)
assert resolver.upload_cache is None

View File

@@ -0,0 +1,171 @@
"""Unit tests for file_store module."""
import uuid
import pytest
from crewai.utilities.file_store import (
clear_files,
clear_task_files,
get_all_files,
get_files,
get_task_files,
store_files,
store_task_files,
)
from crewai.utilities.files import TextFile
class TestFileStore:
"""Tests for synchronous file store operations."""
def setup_method(self) -> None:
"""Set up test fixtures."""
self.crew_id = uuid.uuid4()
self.task_id = uuid.uuid4()
self.test_file = TextFile(source=b"test content")
def teardown_method(self) -> None:
"""Clean up after tests."""
clear_files(self.crew_id)
clear_task_files(self.task_id)
def test_store_and_get_files(self) -> None:
"""Test storing and retrieving crew files."""
files = {"doc": self.test_file}
store_files(self.crew_id, files)
retrieved = get_files(self.crew_id)
assert retrieved is not None
assert "doc" in retrieved
assert retrieved["doc"].read() == b"test content"
def test_get_files_returns_none_when_empty(self) -> None:
"""Test that get_files returns None for non-existent keys."""
new_id = uuid.uuid4()
result = get_files(new_id)
assert result is None
def test_clear_files(self) -> None:
"""Test clearing crew files."""
files = {"doc": self.test_file}
store_files(self.crew_id, files)
clear_files(self.crew_id)
result = get_files(self.crew_id)
assert result is None
def test_store_and_get_task_files(self) -> None:
"""Test storing and retrieving task files."""
files = {"task_doc": self.test_file}
store_task_files(self.task_id, files)
retrieved = get_task_files(self.task_id)
assert retrieved is not None
assert "task_doc" in retrieved
def test_clear_task_files(self) -> None:
"""Test clearing task files."""
files = {"task_doc": self.test_file}
store_task_files(self.task_id, files)
clear_task_files(self.task_id)
result = get_task_files(self.task_id)
assert result is None
def test_get_all_files_merges_crew_and_task(self) -> None:
"""Test that get_all_files merges crew and task files."""
crew_file = TextFile(source=b"crew content")
task_file = TextFile(source=b"task content")
store_files(self.crew_id, {"crew_doc": crew_file})
store_task_files(self.task_id, {"task_doc": task_file})
merged = get_all_files(self.crew_id, self.task_id)
assert merged is not None
assert "crew_doc" in merged
assert "task_doc" in merged
def test_get_all_files_task_overrides_crew(self) -> None:
"""Test that task files override crew files with same name."""
crew_file = TextFile(source=b"crew version")
task_file = TextFile(source=b"task version")
store_files(self.crew_id, {"shared_doc": crew_file})
store_task_files(self.task_id, {"shared_doc": task_file})
merged = get_all_files(self.crew_id, self.task_id)
assert merged is not None
assert merged["shared_doc"].read() == b"task version"
def test_get_all_files_crew_only(self) -> None:
"""Test get_all_files with only crew files."""
store_files(self.crew_id, {"doc": self.test_file})
result = get_all_files(self.crew_id)
assert result is not None
assert "doc" in result
def test_get_all_files_returns_none_when_empty(self) -> None:
"""Test that get_all_files returns None when no files exist."""
new_crew_id = uuid.uuid4()
new_task_id = uuid.uuid4()
result = get_all_files(new_crew_id, new_task_id)
assert result is None
@pytest.mark.asyncio
class TestAsyncFileStore:
"""Tests for asynchronous file store operations."""
async def test_astore_and_aget_files(self) -> None:
"""Test async storing and retrieving crew files."""
from crewai.utilities.file_store import aclear_files, aget_files, astore_files
crew_id = uuid.uuid4()
test_file = TextFile(source=b"async content")
try:
await astore_files(crew_id, {"doc": test_file})
retrieved = await aget_files(crew_id)
assert retrieved is not None
assert "doc" in retrieved
assert retrieved["doc"].read() == b"async content"
finally:
await aclear_files(crew_id)
async def test_aget_all_files(self) -> None:
"""Test async get_all_files merging."""
from crewai.utilities.file_store import (
aclear_files,
aclear_task_files,
aget_all_files,
astore_files,
astore_task_files,
)
crew_id = uuid.uuid4()
task_id = uuid.uuid4()
try:
await astore_files(crew_id, {"crew": TextFile(source=b"crew")})
await astore_task_files(task_id, {"task": TextFile(source=b"task")})
merged = await aget_all_files(crew_id, task_id)
assert merged is not None
assert "crew" in merged
assert "task" in merged
finally:
await aclear_files(crew_id)
await aclear_task_files(task_id)