mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
22 lines
877 B
Python
22 lines
877 B
Python
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from crewai.utilities.embedding_configurator import EmbeddingConfigurator
|
|
|
|
|
|
class TestEmbeddingConfigurator(unittest.TestCase):
|
|
@patch("chromadb.utils.embedding_functions.amazon_bedrock_embedding_function.AmazonBedrockEmbeddingFunction")
|
|
def test_configure_bedrock(self, mock_bedrock_embedder):
|
|
"""Test that the Bedrock embedder is configured correctly."""
|
|
config = {"session": MagicMock()}
|
|
model_name = "amazon.titan-embed-text-v1"
|
|
|
|
embedder = EmbeddingConfigurator()._configure_bedrock(config, model_name)
|
|
|
|
mock_bedrock_embedder.assert_called_once_with(
|
|
session=config["session"],
|
|
model_name=model_name,
|
|
operation_name="InvokeModel",
|
|
)
|
|
self.assertEqual(embedder, mock_bedrock_embedder.return_value)
|