@@ -13,6 +13,9 @@
 # limitations under the License.
 #
 import importlib
+from unittest import mock
+from unittest.mock import patch
+
 from google.api_core import operation as ga_operation
 from agentplatform.preview import rag
 from agentplatform.preview.rag.utils._gapic_utils import (
@@ -34,8 +37,6 @@
     ListRagMetadataResponse,
 )
 from google.cloud import aiplatform
-import mock
-from unittest.mock import patch
 import pytest
 import test_rag_constants_preview
 
 
@@ -15,13 +15,13 @@
 """Tests for vertex_rag.retrieval."""
 
 import importlib
+from unittest import mock
 
 from google.cloud import aiplatform
 from agentplatform import rag
 from google.cloud.aiplatform_v1 import VertexRagServiceAsyncClient
 from google.cloud.aiplatform_v1 import VertexRagServiceClient
 import test_rag_constants as tc
-import mock
 import pytest
 
 
 
@@ -15,12 +15,13 @@
 """Tests for vertex_rag.retrieval_preview."""
 
 import importlib
+from unittest import mock
+
 from google.cloud import aiplatform
 from google.cloud.aiplatform_v1beta1 import VertexRagServiceAsyncClient
 from google.cloud.aiplatform_v1beta1 import VertexRagServiceClient
 import test_rag_constants_preview
 from agentplatform.preview import rag
-import mock
 import pytest
 
 tc = test_rag_constants_preview
 
@@ -13,6 +13,8 @@
 # limitations under the License.
 #
 
+from unittest import mock
+
 from google import auth
 from google.api_core import exceptions
 from google.api_core import operation as ga_operation
@@ -32,7 +34,6 @@
     DeletePersistentResourceRequest,
 )
 import test_constants as tc
-import mock
 import pytest
 
 
 
@@ -16,6 +16,7 @@
 #
 
 import importlib
+from unittest import mock
 
 from google.api_core import exceptions
 from google.api_core import operation
@@ -29,7 +30,6 @@
 import test_constants as tc
 from google.cloud.bigquery import job
 from google.cloud.bigquery_storage_v1.types import stream as gcbqs_stream
-import mock
 import pytest
 import pyarrow as pa
 import ray
 
@@ -15,6 +15,7 @@
 import copy
 import importlib
 import re
+from unittest import mock
 
 from google.api_core import operation as ga_operation
 from google.cloud import aiplatform
@@ -28,7 +29,6 @@
 )
 from google.cloud.aiplatform_v1beta1.types import persistent_resource_service
 import test_constants as tc
-import mock
 import pytest
 
 from google.protobuf import field_mask_pb2  # type: ignore
 
@@ -13,11 +13,11 @@
 # limitations under the License.
 #
 import importlib
+from unittest import mock
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import vertex_ray
 import test_constants as tc
-import mock
 import pytest
 from ray.dashboard.modules import dashboard_sdk as oss_dashboard_sdk
 
 
@@ -19,6 +19,7 @@
 import importlib
 import pickle
 import tempfile
+from unittest import mock
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform.vertex_ray.predict import (
@@ -37,7 +38,6 @@
 import test_constants as tc
 import test_prediction_utils
 
-import mock
 import numpy as np
 import pytest
 import ray
 
@@ -13,11 +13,11 @@
 # limitations under the License.
 #
 import importlib
+from unittest import mock
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import vertex_ray
 import test_constants as tc
-import mock
 import pytest
 import ray
 
 
@@ -20,7 +20,8 @@
 import datetime
 import io
 import json
-import mock
+from unittest import mock
+
 import pytest
 from vertexai.caching import CachedContent
 from google.cloud.aiplatform import initializer
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,8 @@`
`13`	`13`	`# limitations under the License.`
`14`	`14`	`#`
`15`	`15`
	`16`	`+from unittest import mock`
	`17`	`+`
`16`	`18`	`from google import auth`
`17`	`19`	`from google.api_core import exceptions`
`18`	`20`	`from google.api_core import operation as ga_operation`
`@@ -32,7 +34,6 @@`
`32`	`34`	`DeletePersistentResourceRequest,`
`33`	`35`	`)`
`34`	`36`	`import test_constants as tc`
`35`		`-import mock`
`36`	`37`	`import pytest`
`37`	`38`
`38`	`39`