From f79e095419934ae8ee33ff7e3947a5789b9d3b62 Mon Sep 17 00:00:00 2001
From: Matthew Tang <tangmatthew@google.com>
Date: Tue, 5 May 2026 15:51:36 -0700
Subject: [PATCH] chore: Fix flaky unit tests

PiperOrigin-RevId: 910960736
---
 google/cloud/aiplatform/base.py               |  4 +++
 noxfile.py                                    |  2 --
 setup.py                                      |  1 -
 .../preview/feature_store/__init__.py         | 33 +++++++++----------
 .../preview/feature_store/feature_group.py    | 12 +++----
 5 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
index cca3d69064..868ed482d0 100644
--- a/google/cloud/aiplatform/base.py
+++ b/google/cloud/aiplatform/base.py
@@ -465,6 +465,10 @@ class VertexAiResourceNoun(metaclass=abc.ABCMeta):
     service representation of the resource noun.
     """
 
+    def __new__(cls, *args, **kwargs):
+        """Explicit __new__ to prevent argument forwarding to object.__new__."""
+        return super().__new__(cls)
+
     @property
     @classmethod
     @abc.abstractmethod
diff --git a/noxfile.py b/noxfile.py
index ce2d2c3865..7ce394640a 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -62,8 +62,6 @@
     "3.11": ["2.42.0", "2.47.1"],
 }
 UNIT_TEST_STANDARD_DEPENDENCIES = [
-    "mock",
-    "asyncmock",
     "pytest",
     "pytest-cov",
     "pytest-asyncio",
diff --git a/setup.py b/setup.py
index b5ef39736a..a1a80ae409 100644
--- a/setup.py
+++ b/setup.py
@@ -271,7 +271,6 @@
         "kfp >= 2.6.0, < 3.0.0; python_version<'3.13'",
         "pytest-asyncio",
         "pytest-cov",
-        "mock",
         "pytest-xdist",
         "Pillow",
         "scikit-learn<1.6.0; python_version<='3.10'",
diff --git a/vertexai/resources/preview/feature_store/__init__.py b/vertexai/resources/preview/feature_store/__init__.py
index 900b7f38f8..f35be6f0cf 100644
--- a/vertexai/resources/preview/feature_store/__init__.py
+++ b/vertexai/resources/preview/feature_store/__init__.py
@@ -51,21 +51,20 @@
 )
 
 __all__ = (
-    Feature,
-    FeatureGroup,
-    FeatureGroupBigQuerySource,
-    FeatureMonitor,
-    FeatureOnlineStoreType,
-    FeatureOnlineStore,
-    FeatureView,
-    FeatureViewBigQuerySource,
-    FeatureViewReadResponse,
-    FeatureViewVertexRagSource,
-    FeatureViewRegistrySource,
-    IndexConfig,
-    IndexConfig,
-    TreeAhConfig,
-    BruteForceConfig,
-    DistanceMeasureType,
-    AlgorithmConfig,
+    "Feature",
+    "FeatureGroup",
+    "FeatureGroupBigQuerySource",
+    "FeatureMonitor",
+    "FeatureOnlineStoreType",
+    "FeatureOnlineStore",
+    "FeatureView",
+    "FeatureViewBigQuerySource",
+    "FeatureViewReadResponse",
+    "FeatureViewVertexRagSource",
+    "FeatureViewRegistrySource",
+    "IndexConfig",
+    "TreeAhConfig",
+    "BruteForceConfig",
+    "DistanceMeasureType",
+    "AlgorithmConfig",
 )
diff --git a/vertexai/resources/preview/feature_store/feature_group.py b/vertexai/resources/preview/feature_store/feature_group.py
index 9d04339399..42db71bf4e 100644
--- a/vertexai/resources/preview/feature_store/feature_group.py
+++ b/vertexai/resources/preview/feature_store/feature_group.py
@@ -20,20 +20,20 @@
 from google.cloud.aiplatform import base, initializer
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform.compat.types import (
-    feature as gca_feature,
     feature_group as gca_feature_group,
-    io as gca_io,
     feature_monitor_v1beta1 as gca_feature_monitor,
+    feature as gca_feature,
+    io as gca_io,
 )
-from vertexai.resources.preview.feature_store.utils import (
-    FeatureGroupBigQuerySource,
-)
-from vertexai.resources.preview.feature_store import (
+from vertexai.resources.preview.feature_store.feature import (
     Feature,
 )
 from vertexai.resources.preview.feature_store.feature_monitor import (
     FeatureMonitor,
 )
+from vertexai.resources.preview.feature_store.utils import (
+    FeatureGroupBigQuerySource,
+)
 
 
 _LOGGER = base.Logger(__name__)