include unit tests

Author: AnFreTh

tests/test_base.py

@@ -0,0 +1,155 @@
+import pytest
+import inspect
+import torch
+import os
+import importlib
+from mambular.base_models.basemodel import BaseModel
+
+# Paths for models and configs
+MODEL_MODULE_PATH = "mambular.base_models"
+CONFIG_MODULE_PATH = "mambular.configs"
+
+# Discover all models
+model_classes = []
+for filename in os.listdir(os.path.dirname(__file__) + "/../mambular/base_models"):
+    if filename.endswith(".py") and filename not in [
+        "__init__.py",
+        "basemodel.py",
+        "lightning_wrapper.py",
+        "bayesian_tabm.py",
+    ]:
+        module_name = f"{MODEL_MODULE_PATH}.{filename[:-3]}"
+        module = importlib.import_module(module_name)
+
+        for name, obj in inspect.getmembers(module, inspect.isclass):
+            if issubclass(obj, BaseModel) and obj is not BaseModel:
+                model_classes.append(obj)
+
+
+def get_model_config(model_class):
+    """Dynamically load the correct config class for each model."""
+    model_name = model_class.__name__  # e.g., "Mambular"
+    config_class_name = f"Default{model_name}Config"  # e.g., "DefaultMambularConfig"
+
+    try:
+        config_module = importlib.import_module(
+            f"{CONFIG_MODULE_PATH}.{model_name.lower()}_config"
+        )
+        config_class = getattr(config_module, config_class_name)
+        return config_class()  # Instantiate config
+    except (ModuleNotFoundError, AttributeError) as e:
+        pytest.fail(
+            f"Could not find or instantiate config {config_class_name} for {model_name}: {e}"
+        )
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_inherits_base_model(model_class):
+    """Test that each model correctly inherits from BaseModel."""
+    assert issubclass(
+        model_class, BaseModel
+    ), f"{model_class.__name__} should inherit from BaseModel."
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_has_forward_method(model_class):
+    """Test that each model has a forward method with *data."""
+    assert hasattr(
+        model_class, "forward"
+    ), f"{model_class.__name__} is missing a forward method."
+
+    sig = inspect.signature(model_class.forward)
+    assert any(
+        p.kind == inspect.Parameter.VAR_POSITIONAL for p in sig.parameters.values()
+    ), f"{model_class.__name__}.forward should have *data argument."
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_takes_config(model_class):
+    """Test that each model accepts a config argument."""
+    sig = inspect.signature(model_class.__init__)
+    assert (
+        "config" in sig.parameters
+    ), f"{model_class.__name__} should accept a 'config' parameter."
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_has_num_classes(model_class):
+    """Test that each model accepts a num_classes argument."""
+    sig = inspect.signature(model_class.__init__)
+    assert (
+        "num_classes" in sig.parameters
+    ), f"{model_class.__name__} should accept a 'num_classes' parameter."
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_calls_super_init(model_class):
+    """Test that each model calls super().__init__(config=config, **kwargs)."""
+    source = inspect.getsource(model_class.__init__)
+    assert (
+        "super().__init__(config=config" in source
+    ), f"{model_class.__name__} should call super().__init__(config=config, **kwargs)."
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_initialization(model_class):
+    """Test that each model can be initialized with its correct config."""
+    config = get_model_config(model_class)
+    feature_info = (
+        {
+            "A": {
+                "preprocessing": "imputer -> check_positive -> box-cox",
+                "dimension": 1,
+                "categories": None,
+            }
+        },
+        {
+            "sibsp": {
+                "preprocessing": "imputer -> continuous_ordinal",
+                "dimension": 1,
+                "categories": 8,
+            }
+        },
+        {},
+    )  # Mock feature info
+
+    try:
+        model = model_class(
+            feature_information=feature_info, num_classes=3, config=config
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to initialize {model_class.__name__}: {e}")
+
+
+@pytest.mark.parametrize("model_class", model_classes)
+def test_model_defines_key_attributes(model_class):
+    """Test that each model defines expected attributes like returns_ensemble"""
+    config = get_model_config(model_class)
+    feature_info = (
+        {
+            "A": {
+                "preprocessing": "imputer -> check_positive -> box-cox",
+                "dimension": 1,
+                "categories": None,
+            }
+        },
+        {
+            "sibsp": {
+                "preprocessing": "imputer -> continuous_ordinal",
+                "dimension": 1,
+                "categories": 8,
+            }
+        },
+        {},
+    )  # Mock feature info
+
+    try:
+        model = model_class(
+            feature_information=feature_info, num_classes=3, config=config
+        )
+    except TypeError as e:
+        pytest.fail(f"Failed to initialize {model_class.__name__}: {e}")
+
+    expected_attrs = ["returns_ensemble"]
+    for attr in expected_attrs:
+        assert hasattr(model, attr), f"{model_class.__name__} should define '{attr}'."

tests/test_classifier.py

@@ -0,0 +1,115 @@
+import pytest
+import inspect
+import importlib
+import os
+import dataclasses
+import typing
+from mambular.configs.base_config import BaseConfig  # Ensure correct path
+
+CONFIG_MODULE_PATH = "mambular.configs"
+config_classes = []
+
+# Discover all config classes in mambular/configs/
+for filename in os.listdir(os.path.dirname(__file__) + "/../mambular/configs"):
+    if (
+        filename.endswith(".py")
+        and filename != "base_config.py"
+        and not filename.startswith("__")
+    ):
+        module_name = f"{CONFIG_MODULE_PATH}.{filename[:-3]}"
+        module = importlib.import_module(module_name)
+
+        for name, obj in inspect.getmembers(module, inspect.isclass):
+            if issubclass(obj, BaseConfig) and obj is not BaseConfig:
+                config_classes.append(obj)
+
+
+@pytest.mark.parametrize("config_class", config_classes)
+def test_config_inherits_baseconfig(config_class):
+    """Test that each config class correctly inherits from BaseConfig."""
+    assert issubclass(
+        config_class, BaseConfig
+    ), f"{config_class.__name__} should inherit from BaseConfig."
+
+
+@pytest.mark.parametrize("config_class", config_classes)
+def test_config_instantiation(config_class):
+    """Test that each config class can be instantiated without errors."""
+    try:
+        config = config_class()
+    except Exception as e:
+        pytest.fail(f"Failed to instantiate {config_class.__name__}: {e}")
+
+
+@pytest.mark.parametrize("config_class", config_classes)
+def test_config_has_expected_attributes(config_class):
+    """Test that each config has all required attributes from BaseConfig."""
+    base_attrs = {field.name for field in dataclasses.fields(BaseConfig)}
+    config_attrs = {field.name for field in dataclasses.fields(config_class)}
+
+    missing_attrs = base_attrs - config_attrs
+    assert (
+        not missing_attrs
+    ), f"{config_class.__name__} is missing attributes: {missing_attrs}"
+
+
+@pytest.mark.parametrize("config_class", config_classes)
+def test_config_default_values(config_class):
+    """Ensure that each config class has default values assigned correctly."""
+    config = config_class()
+
+    for field in dataclasses.fields(config_class):
+        attr = field.name
+        expected_type = field.type
+
+        assert hasattr(
+            config, attr
+        ), f"{config_class.__name__} is missing attribute '{attr}'."
+
+        value = getattr(config, attr)
+
+        # Handle generic types properly
+        origin = typing.get_origin(expected_type)
+
+        if origin is typing.Literal:
+            # If the field is a Literal, ensure the value is one of the allowed options
+            allowed_values = typing.get_args(expected_type)
+            assert (
+                value in allowed_values
+            ), f"{config_class.__name__}.{attr} has incorrect value: expected one of {allowed_values}, got {value}"
+        elif origin is typing.Union:
+            # For Union types (e.g., Optional[str]), check if value matches any type in the union
+            allowed_types = typing.get_args(expected_type)
+            assert any(
+                isinstance(value, t) for t in allowed_types
+            ), f"{config_class.__name__}.{attr} has incorrect type: expected one of {allowed_types}, got {type(value)}"
+        elif origin is not None:
+            # If it's another generic type (e.g., list[str]), check against the base type
+            assert (
+                isinstance(value, origin) or value is None
+            ), f"{config_class.__name__}.{attr} has incorrect type: expected {expected_type}, got {type(value)}"
+        else:
+            # Standard type check
+            assert (
+                isinstance(value, expected_type) or value is None
+            ), f"{config_class.__name__}.{attr} has incorrect type: expected {expected_type}, got {type(value)}"
+
+
+@pytest.mark.parametrize("config_class", config_classes)
+def test_config_allows_updates(config_class):
+    """Ensure that config values can be updated and remain type-consistent."""
+    config = config_class()
+
+    update_values = {
+        "lr": 0.01,
+        "d_model": 128,
+        "embedding_type": "plr",
+        "activation": lambda x: x,  # Function update
+    }
+
+    for attr, new_value in update_values.items():
+        if hasattr(config, attr):
+            setattr(config, attr, new_value)
+            assert (
+                getattr(config, attr) == new_value
+            ), f"{config_class.__name__}.{attr} did not update correctly."