"""测试 LightGBM 模型 验证 LightGBMModel 的训练、预测、保存和加载功能。 """ import os import tempfile import numpy as np import polars as pl import pytest from src.training.components.models.lightgbm import LightGBMModel class TestLightGBMModel: """LightGBMModel 测试类""" def test_init_default(self): """测试默认初始化""" model = LightGBMModel() assert model.name == "lightgbm" assert model.params["objective"] == "regression" assert model.params["metric"] == "rmse" assert model.params["num_leaves"] == 31 assert model.params["learning_rate"] == 0.05 assert model.n_estimators == 100 assert model.model is None def test_init_custom(self): """测试自定义参数""" model = LightGBMModel( objective="huber", metric="mae", num_leaves=50, learning_rate=0.1, n_estimators=200, ) assert model.params["objective"] == "huber" assert model.params["metric"] == "mae" assert model.params["num_leaves"] == 50 assert model.params["learning_rate"] == 0.1 assert model.n_estimators == 200 def test_fit_success(self): """测试正常训练""" # 创建简单回归数据 X = pl.DataFrame( { "feature1": [1.0, 2.0, 3.0, 4.0, 5.0], "feature2": [2.0, 4.0, 6.0, 8.0, 10.0], } ) y = pl.Series("target", [1.5, 3.0, 4.5, 6.0, 7.5]) model = LightGBMModel(n_estimators=10) result = model.fit(X, y) # 验证返回 self（支持链式调用） assert result is model # 验证模型已训练 assert model.model is not None # 验证特征名称已保存 assert model.feature_names_ == ["feature1", "feature2"] def test_predict_before_fit(self): """测试未训练就预测""" X = pl.DataFrame( { "feature1": [1.0, 2.0], "feature2": [2.0, 4.0], } ) model = LightGBMModel() with pytest.raises(RuntimeError, match="模型尚未训练"): model.predict(X) def test_predict_success(self): """测试正常预测""" # 创建回归数据 np.random.seed(42) n_samples = 100 X_train = pl.DataFrame( { "feature1": np.random.randn(n_samples), "feature2": np.random.randn(n_samples), } ) # y = 2*feature1 + 3*feature2 + noise y_train = pl.Series( "target", 2 * X_train["feature1"] + 3 * X_train["feature2"] + np.random.randn(n_samples) * 0.1, ) model = LightGBMModel(n_estimators=20, learning_rate=0.1) model.fit(X_train, y_train) # 预测新数据（使用明显不同的值） X_test = pl.DataFrame( { "feature1": [-2.0, 3.0], "feature2": [-1.0, 4.0], } ) predictions = model.predict(X_test) # 验证预测结果格式 assert isinstance(predictions, np.ndarray) assert len(predictions) == 2 # 验证预测值是数值 assert all(np.isfinite(predictions)) # 验证单调性（第二个样本的 feature 值更大，预测值也应更大） assert predictions[1] > predictions[0] def test_feature_importance_before_fit(self): """测试未训练就获取特征重要性""" model = LightGBMModel() assert model.feature_importance() is None def test_feature_importance_after_fit(self): """测试训练后获取特征重要性""" X = pl.DataFrame( { "feature1": np.random.randn(100), "feature2": np.random.randn(100), } ) y = pl.Series("target", X["feature1"] * 2 + X["feature2"] * 0.1) model = LightGBMModel(n_estimators=10) model.fit(X, y) importance = model.feature_importance() # 验证特征重要性格式 assert importance is not None assert len(importance) == 2 assert "feature1" in importance.index assert "feature2" in importance.index # feature1 的系数更大，重要性应该更高 assert importance["feature1"] >= importance["feature2"] def test_save_before_fit(self): """测试未训练就保存""" model = LightGBMModel() with pytest.raises(RuntimeError, match="模型尚未训练"): model.save("dummy.txt") def test_save_and_load(self): """测试保存和加载""" # 训练模型 X = pl.DataFrame( { "feature1": [1.0, 2.0, 3.0, 4.0, 5.0], "feature2": [2.0, 4.0, 6.0, 8.0, 10.0], } ) y = pl.Series("target", [2.0, 4.0, 6.0, 8.0, 10.0]) model = LightGBMModel(n_estimators=10, learning_rate=0.1) model.fit(X, y) # 保存前预测 X_test = pl.DataFrame( { "feature1": [6.0], "feature2": [12.0], } ) pred_before = model.predict(X_test) # 保存到临时文件 with tempfile.TemporaryDirectory() as tmpdir: save_path = os.path.join(tmpdir, "model.txt") model.save(save_path) # 加载模型 loaded_model = LightGBMModel.load(save_path) # 验证加载后预测结果相同 pred_after = loaded_model.predict(X_test) assert pred_after[0] == pytest.approx(pred_before[0], rel=1e-5) # 验证元数据已恢复 assert loaded_model.feature_names_ == ["feature1", "feature2"] assert loaded_model.n_estimators == 10 def test_registration(self): """测试模型已注册到 registry""" from src.training.registry import ModelRegistry # 重新导入模型模块以确保注册（处理其他测试 clear 注册表的情况） import importlib import src.training.components.models.lightgbm as lightgbm_module importlib.reload(lightgbm_module) from src.training.components.models.lightgbm import ( LightGBMModel as ReloadedModel, ) model_class = ModelRegistry.get_model("lightgbm") assert model_class is ReloadedModel def test_fit_predict_consistency(self): """测试多次预测结果一致""" X = pl.DataFrame( { "feature1": np.random.randn(50), "feature2": np.random.randn(50), } ) y = pl.Series("target", X["feature1"] + X["feature2"]) model = LightGBMModel(n_estimators=10) model.fit(X, y) X_test = pl.DataFrame( { "feature1": [1.0, 2.0, 3.0], "feature2": [1.0, 2.0, 3.0], } ) # 多次预测应该返回相同结果 pred1 = model.predict(X_test) pred2 = model.predict(X_test) np.testing.assert_array_almost_equal(pred1, pred2) if __name__ == "__main__": pytest.main([__file__, "-v"])