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.
87 lines
2.8 KiB
87 lines
2.8 KiB
5 months ago
|
from unittest.mock import Mock, patch
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from swarms.models import TimmModel
|
||
|
|
|
||
|
|
|
||
|
|
def test_timm_model_init():
|
||
|
|
with patch("swarms.models.timm.list_models") as mock_list_models:
|
||
|
|
model_name = "resnet18"
|
||
|
|
pretrained = True
|
||
|
|
in_chans = 3
|
||
|
|
timm_model = TimmModel(model_name, pretrained, in_chans)
|
||
|
|
mock_list_models.assert_called_once()
|
||
|
|
assert timm_model.model_name == model_name
|
||
|
|
assert timm_model.pretrained == pretrained
|
||
|
|
assert timm_model.in_chans == in_chans
|
||
|
|
assert timm_model.models == mock_list_models.return_value
|
||
|
|
|
||
|
|
|
||
|
|
def test_timm_model_call():
|
||
|
|
with patch("swarms.models.timm.create_model") as mock_create_model:
|
||
|
|
model_name = "resnet18"
|
||
|
|
pretrained = True
|
||
|
|
in_chans = 3
|
||
|
|
timm_model = TimmModel(model_name, pretrained, in_chans)
|
||
|
|
task = torch.rand(1, in_chans, 224, 224)
|
||
|
|
result = timm_model(task)
|
||
|
|
mock_create_model.assert_called_once_with(
|
||
|
|
model_name, pretrained=pretrained, in_chans=in_chans
|
||
|
|
)
|
||
|
|
assert result == mock_create_model.return_value(task)
|
||
|
|
|
||
|
|
|
||
|
|
def test_timm_model_list_models():
|
||
|
|
with patch("swarms.models.timm.list_models") as mock_list_models:
|
||
|
|
model_name = "resnet18"
|
||
|
|
pretrained = True
|
||
|
|
in_chans = 3
|
||
|
|
timm_model = TimmModel(model_name, pretrained, in_chans)
|
||
|
|
result = timm_model.list_models()
|
||
|
|
mock_list_models.assert_called_once()
|
||
|
|
assert result == mock_list_models.return_value
|
||
|
|
|
||
|
|
|
||
|
|
def test_get_supported_models():
|
||
|
|
model_handler = TimmModel()
|
||
|
|
supported_models = model_handler._get_supported_models()
|
||
|
|
assert isinstance(supported_models, list)
|
||
|
|
assert len(supported_models) > 0
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_model(sample_model_info):
|
||
|
|
model_handler = TimmModel()
|
||
|
|
model = model_handler._create_model(sample_model_info)
|
||
|
|
assert isinstance(model, torch.nn.Module)
|
||
|
|
|
||
|
|
|
||
|
|
def test_call(sample_model_info):
|
||
|
|
model_handler = TimmModel()
|
||
|
|
input_tensor = torch.randn(1, 3, 224, 224)
|
||
|
|
output_shape = model_handler.__call__(sample_model_info, input_tensor)
|
||
|
|
assert isinstance(output_shape, torch.Size)
|
||
|
|
|
||
|
|
|
||
|
|
def test_get_supported_models_mock():
|
||
|
|
model_handler = TimmModel()
|
||
|
|
model_handler._get_supported_models = Mock(
|
||
|
|
return_value=["resnet18", "resnet50"]
|
||
|
|
)
|
||
|
|
supported_models = model_handler._get_supported_models()
|
||
|
|
assert supported_models == ["resnet18", "resnet50"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_model_mock(sample_model_info):
|
||
|
|
model_handler = TimmModel()
|
||
|
|
model_handler._create_model = Mock(return_value=torch.nn.Module())
|
||
|
|
model = model_handler._create_model(sample_model_info)
|
||
|
|
assert isinstance(model, torch.nn.Module)
|
||
|
|
|
||
|
|
|
||
|
|
def test_coverage_report():
|
||
|
|
# Install pytest-cov
|
||
|
|
# Run tests with coverage report
|
||
|
|
pytest.main(["--cov=my_module", "--cov-report=html"])
|