You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
34 lines
1.1 KiB
34 lines
1.1 KiB
1 year ago
|
import pytest
|
||
|
|
from unittest.mock import Mock, patch
|
||
|
|
from swarms.embeddings.pegasus import PegasusEmbedding
|
||
|
|
|
||
|
|
|
||
|
|
def test_init():
|
||
|
|
with patch("your_module.Pegasus") as MockPegasus:
|
||
|
|
embedder = PegasusEmbedding(modality="text")
|
||
|
|
MockPegasus.assert_called_once()
|
||
|
|
assert embedder.pegasus == MockPegasus.return_value
|
||
|
|
|
||
|
|
|
||
|
|
def test_init_exception():
|
||
|
|
with patch("your_module.Pegasus", side_effect=Exception("Test exception")):
|
||
|
|
with pytest.raises(Exception) as e:
|
||
|
|
embedder = PegasusEmbedding(modality="text")
|
||
|
|
assert str(e.value) == "Test exception"
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed():
|
||
|
|
with patch("your_module.Pegasus") as MockPegasus:
|
||
|
|
embedder = PegasusEmbedding(modality="text")
|
||
|
|
embedder.embed("Hello world")
|
||
|
|
MockPegasus.return_value.embed.assert_called_once()
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed_exception():
|
||
|
|
with patch("your_module.Pegasus") as MockPegasus:
|
||
|
|
MockPegasus.return_value.embed.side_effect = Exception("Test exception")
|
||
|
|
embedder = PegasusEmbedding(modality="text")
|
||
|
|
with pytest.raises(Exception) as e:
|
||
|
|
embedder.embed("Hello world")
|
||
|
|
assert str(e.value) == "Test exception"
|