From 358a1d735f8909ad54289a2a62c2ca39039d7686 Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Tue, 20 Jan 2026 11:50:03 -0800 Subject: [PATCH] chore: fix broken test PiperOrigin-RevId: 858688708 --- tests/unit/aiplatform/test_metadata_models.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/aiplatform/test_metadata_models.py b/tests/unit/aiplatform/test_metadata_models.py index b7120c2310..6d0fedcda8 100644 --- a/tests/unit/aiplatform/test_metadata_models.py +++ b/tests/unit/aiplatform/test_metadata_models.py @@ -18,6 +18,7 @@ import datetime import os import pickle +import sys from importlib import reload from unittest import mock from unittest.mock import patch @@ -41,6 +42,7 @@ from sklearn.linear_model import LinearRegression import tensorflow as tf import xgboost as xgb +import xgboost.sklearn # project @@ -339,6 +341,8 @@ def setup_method(self): reload(initializer) reload(metadata) reload(aiplatform) + reload(xgboost.sklearn) + reload(xgb) def teardown_method(self): initializer.global_pool.shutdown(wait=True)