26
26
from google .cloud import aiplatform
27
27
from google .cloud .aiplatform import base
28
28
from google .cloud .aiplatform import initializer
29
- from google .cloud .aiplatform import models
30
29
31
30
from google .cloud .aiplatform .compat .services import (
32
31
deployment_resource_pool_service_client ,
@@ -232,7 +231,7 @@ def test_constructor_gets_drp(self, get_drp_mock):
232
231
location = _TEST_LOCATION ,
233
232
credentials = _TEST_CREDENTIALS ,
234
233
)
235
- models .DeploymentResourcePool (_TEST_DRP_NAME )
234
+ aiplatform .DeploymentResourcePool (_TEST_DRP_NAME )
236
235
get_drp_mock .assert_called_once_with (
237
236
name = _TEST_DRP_NAME , retry = base ._DEFAULT_RETRY
238
237
)
@@ -242,7 +241,7 @@ def test_constructor_with_conflicting_location_fails(self):
242
241
"""Passing a full resource name with `_TEST_LOCATION` and providing `_TEST_LOCATION_2` as location"""
243
242
244
243
with pytest .raises (RuntimeError ) as err :
245
- models .DeploymentResourcePool (_TEST_DRP_NAME , location = _TEST_LOCATION_2 )
244
+ aiplatform .DeploymentResourcePool (_TEST_DRP_NAME , location = _TEST_LOCATION_2 )
246
245
247
246
assert err .match (
248
247
regexp = r"is provided, but different from the resource location"
@@ -251,7 +250,7 @@ def test_constructor_with_conflicting_location_fails(self):
251
250
@pytest .mark .usefixtures ("create_drp_mock" )
252
251
@pytest .mark .parametrize ("sync" , [True , False ])
253
252
def test_create (self , create_drp_mock , sync ):
254
- test_drp = models .DeploymentResourcePool .create (
253
+ test_drp = aiplatform .DeploymentResourcePool .create (
255
254
deployment_resource_pool_id = _TEST_ID ,
256
255
machine_type = _TEST_MACHINE_TYPE ,
257
256
min_replica_count = 10 ,
@@ -285,7 +284,7 @@ def test_create(self, create_drp_mock, sync):
285
284
@pytest .mark .usefixtures ("create_drp_mock" )
286
285
@pytest .mark .parametrize ("sync" , [True , False ])
287
286
def test_create_with_timeout (self , create_drp_mock , sync ):
288
- test_drp = models .DeploymentResourcePool .create (
287
+ test_drp = aiplatform .DeploymentResourcePool .create (
289
288
deployment_resource_pool_id = _TEST_ID ,
290
289
machine_type = _TEST_MACHINE_TYPE ,
291
290
min_replica_count = 10 ,
@@ -319,17 +318,17 @@ def test_create_with_timeout(self, create_drp_mock, sync):
319
318
320
319
@pytest .mark .usefixtures ("list_drp_mock" )
321
320
def test_list (self , list_drp_mock ):
322
- drp_list = models .DeploymentResourcePool .list ()
321
+ drp_list = aiplatform .DeploymentResourcePool .list ()
323
322
324
323
list_drp_mock .assert_called_once ()
325
324
326
325
for drp in drp_list :
327
- assert isinstance (drp , models .DeploymentResourcePool )
326
+ assert isinstance (drp , aiplatform .DeploymentResourcePool )
328
327
329
328
@pytest .mark .usefixtures ("delete_drp_mock" , "get_drp_mock" )
330
329
@pytest .mark .parametrize ("sync" , [True , False ])
331
330
def test_delete (self , delete_drp_mock , get_drp_mock , sync ):
332
- test_drp = models .DeploymentResourcePool (
331
+ test_drp = aiplatform .DeploymentResourcePool (
333
332
deployment_resource_pool_name = _TEST_DRP_NAME
334
333
)
335
334
test_drp .delete (sync = sync )
@@ -341,7 +340,7 @@ def test_delete(self, delete_drp_mock, get_drp_mock, sync):
341
340
342
341
@pytest .mark .usefixtures ("query_deployed_models_mock" , "get_drp_mock" )
343
342
def test_query_deployed_models (self , query_deployed_models_mock , get_drp_mock ):
344
- test_drp = models .DeploymentResourcePool (
343
+ test_drp = aiplatform .DeploymentResourcePool (
345
344
deployment_resource_pool_name = _TEST_DRP_NAME
346
345
)
347
346
dm_refs = test_drp .query_deployed_models ()
0 commit comments