adding a sample test

sklearn_genetic/tests/test_genetic_search.py

@@ -10,6 +10,7 @@
 from sklearn.metrics import make_scorer
 
 import numpy as np
+import os
 
 from .. import GASearchCV
 from ..space import Integer, Categorical, Continuous
@@ -19,6 +20,7 @@
     ConsecutiveStopping,
     TimerStopping,
     ProgressBar,
+    ModelCheckpoint
 )
 
 from ..schedules import ExponentialAdapter, InverseAdapter
@@ -659,3 +661,51 @@ def test_expected_ga_schedulers():
     assert "params" in cv_result_keys
 
     assert crossover_scheduler.current_value + mutation_scheduler.current_value <= 1
+
+
+def test_checkpoint_functionality():
+    clf = SGDClassifier(loss="modified_huber", fit_intercept=True)
+    gen = 5
+    evolved_estimator = GASearchCV(
+        clf,
+        cv=3,
+        scoring="accuracy",
+        population_size=6,
+        generations=gen,
+        tournament_size=3,
+        param_grid={
+            "l1_ratio": Continuous(0, 1),
+            "alpha": Continuous(1e-4, 1),
+            "average": Categorical([True, False]),
+        },
+    )
+    checkpoint_path = 'test_checkpoint.pkl'
+    checkpoint = ModelCheckpoint(checkpoint_path=checkpoint_path)  # noqa
+    evolved_estimator.fit(X_train, y_train, callbacks=checkpoint)
+
+    checkpoint_data = checkpoint.load()
+
+    assert 'estimator' in checkpoint_data['estimator_state']
+    assert 'algorithm' in checkpoint_data['estimator_state']
+    assert 'logbook' in checkpoint_data
+
+    restored_estimator = GASearchCV(**checkpoint_data['estimator_state'])
+
+    assert restored_estimator.algorithm == checkpoint_data['estimator_state']['algorithm']  # noqa
+
+    assert len(checkpoint_data['logbook']) == gen + 1
+
+    test_estimator = GASearchCV(
+        clf,
+        param_grid={
+            "l1_ratio": Continuous(0, 1),
+            "alpha": Continuous(1e-1, 1),
+            "average": Categorical([True, False]),
+    },)
+
+    test_estimator.load('checkpoint_path')
+
+    assert restored_estimator.algorithm == test_estimator.algorithm  # noqa
+
+    if os.path.exists(checkpoint_path):
+        os.remove(checkpoint_path)