52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""Unit tests for models __init__.py module."""
|
|
|
|
from model_manager.utils.models import (
|
|
ExperimentStatus,
|
|
TrainModelParams,
|
|
TrainModelResult,
|
|
)
|
|
|
|
|
|
def test_experiment_status_import():
|
|
"""Test that ExperimentStatus can be imported from models package."""
|
|
assert ExperimentStatus is not None
|
|
assert hasattr(ExperimentStatus, 'ORCHESTRATOR_WAITING_PROC')
|
|
assert hasattr(ExperimentStatus, 'TRAINING_SUCCESS')
|
|
|
|
|
|
def test_train_model_params_import():
|
|
"""Test that TrainModelParams can be imported from models package."""
|
|
assert TrainModelParams is not None
|
|
assert callable(TrainModelParams)
|
|
|
|
|
|
def test_train_model_result_import():
|
|
"""Test that TrainModelResult can be imported from models package."""
|
|
assert TrainModelResult is not None
|
|
# Dataclasses have __dataclass_fields__
|
|
assert hasattr(TrainModelResult, '__dataclass_fields__')
|
|
|
|
|
|
def test_all_exports():
|
|
"""Test that __all__ contains all expected exports."""
|
|
from model_manager.utils.models import __all__
|
|
|
|
assert 'ExperimentStatus' in __all__
|
|
assert 'TrainModelParams' in __all__
|
|
assert 'TrainModelResult' in __all__
|
|
assert len(__all__) == 3
|
|
|
|
|
|
def test_no_extra_exports():
|
|
"""Test that only expected items are exported."""
|
|
import model_manager.utils.models as models_module
|
|
|
|
# Get all public attributes (not starting with _)
|
|
public_attrs = [attr for attr in dir(models_module) if not attr.startswith('_')]
|
|
|
|
# Should only have the 3 main classes
|
|
expected_public = {'ExperimentStatus', 'TrainModelParams', 'TrainModelResult'}
|
|
|
|
# Check that our expected classes are present
|
|
assert expected_public.issubset(set(public_attrs))
|