diff --git a/tests/unit/code_test.py b/tests/unit/code_test.py
index 6b6d60ac..3ae1c2dd 100644
--- a/tests/unit/code_test.py
+++ b/tests/unit/code_test.py
@@ -1,4 +1,5 @@
 import numpy as np
+import lightgbm
 from azureml.core.run import Run
 from unittest.mock import Mock
 from diabetes_regression.training.train import train_model