update: integration tests

Author: morgandu

tests/system/aiplatform/e2e_base.py

@@ -58,12 +58,12 @@ def setup_method(self):
  importlib.reload(initializer)
  importlib.reload(aiplatform)
 
- @pytest.fixture()
+ @pytest.fixture(scope="class")
  def shared_state(self) -> Generator[Dict[str, Any], None, None]:
  shared_state = {}
  yield shared_state
 
- @pytest.fixture()
+ @pytest.fixture(scope="class")
  def prepare_staging_bucket(
  self, shared_state: Dict[str, Any]
  ) -> Generator[storage.bucket.Bucket, None, None]:
@@ -80,7 +80,7 @@ def prepare_staging_bucket(
  )
  yield
 
- @pytest.fixture()
+ @pytest.fixture(scope="class")
  def delete_staging_bucket(self, shared_state: Dict[str, Any]):
  """Delete the staging bucket and all it's contents"""
 
@@ -90,7 +90,7 @@ def delete_staging_bucket(self, shared_state: Dict[str, Any]):
  bucket = shared_state["bucket"]
  bucket.delete(force=True)
 
- @pytest.fixture(autouse=True)
+ @pytest.fixture(scope="class", autouse=True)
  def teardown(self, shared_state: Dict[str, Any]):
  """Delete every Vertex AI resource created during test"""
 
@@ -104,9 +104,7 @@ def teardown(self, shared_state: Dict[str, Any]):
 
  for resource in shared_state["resources"]:
  try:
- if isinstance(
- resource, (aiplatform.Endpoint or aiplatform.Featurestore)
- ):
+ if isinstance(resource, (aiplatform.Endpoint, aiplatform.Featurestore)):
  # For endpoint, undeploy model then delete endpoint
  # For featurestore, force delete its entity_types and features with the featurestore
  resource.delete(force=True)

tests/system/aiplatform/test_featurestore.py

@@ -44,80 +44,101 @@ class TestFeaturestore(e2e_base.TestEndToEnd):
 
  _temp_prefix = "temp_vertex_sdk_e2e_featurestore_test"
 
- def test_end_to_end(self, shared_state, caplog):
-
- caplog.set_level(logging.INFO)
-
+ def test_create_get_list_featurestore(self, shared_state):
  aiplatform.init(
  project=e2e_base._PROJECT, location=e2e_base._LOCATION,
  )
 
  base_list_featurestores = len(aiplatform.Featurestore.list())
- base_list_searched_features = len(aiplatform.Feature.search())
+ shared_state["base_list_searched_features"] = len(aiplatform.Feature.search())
 
- # Featurestore
  featurestore_id = self._make_display_name(key=_TEST_FEATURESTORE_ID).replace(
  "-", "_"
  )[:60]
  featurestore = aiplatform.Featurestore.create(featurestore_id=featurestore_id)
+
  shared_state["resources"] = [featurestore]
+ shared_state["featurestore"] = featurestore
+ shared_state["featurestore_name"] = featurestore.resource_name
 
  get_featurestore = aiplatform.Featurestore(
  featurestore_name=featurestore.resource_name
  )
  assert featurestore.resource_name == get_featurestore.resource_name
 
  list_featurestores = aiplatform.Featurestore.list()
- assert len(list_featurestores) - base_list_featurestores == 1
+ assert (len(list_featurestores) - base_list_featurestores) == 1
+
+ def test_create_get_list_entity_types(self, shared_state):
+
+ assert shared_state["featurestore"]
+ assert shared_state["featurestore_name"]
+
+ featurestore = shared_state["featurestore"]
+ featurestore_name = shared_state["featurestore_name"]
 
- # EntityType
+ aiplatform.init(
+ project=e2e_base._PROJECT, location=e2e_base._LOCATION,
+ )
 
- # User EntityType
- user_entity_type_id = _TEST_USER_ENTITY_TYPE_ID
+ # Users
  user_entity_type = featurestore.create_entity_type(
- entity_type_id=user_entity_type_id
+ entity_type_id=_TEST_USER_ENTITY_TYPE_ID
  )
+ shared_state["user_entity_type"] = user_entity_type
+ shared_state["user_entity_type_name"] = user_entity_type.resource_name
 
  get_user_entity_type = featurestore.get_entity_type(
- entity_type_id=user_entity_type_id
+ entity_type_id=_TEST_USER_ENTITY_TYPE_ID
  )
  assert user_entity_type.resource_name == get_user_entity_type.resource_name
 
- # Movie EntityType
- movie_entity_type_id = _TEST_MOVIE_ENTITY_TYPE_ID
+ # Movies
  movie_entity_type = aiplatform.EntityType.create(
- entity_type_id=movie_entity_type_id,
- featurestore_name=featurestore.resource_name,
+ entity_type_id=_TEST_MOVIE_ENTITY_TYPE_ID,
+ featurestore_name=featurestore_name,
  )
+ shared_state["movie_entity_type"] = movie_entity_type
+ shared_state["movie_entity_type_name"] = movie_entity_type.resource_name
 
  get_movie_entity_type = aiplatform.EntityType(
  entity_type_name=movie_entity_type.resource_name
  )
  assert movie_entity_type.resource_name == get_movie_entity_type.resource_name
 
  list_entity_types = aiplatform.EntityType.list(
- featurestore_name=featurestore.resource_name
+ featurestore_name=featurestore_name
  )
  assert len(list_entity_types) == 2
 
- # Feature
+ def test_create_batch_create_get_list_features(self, shared_state):
+
+ assert shared_state["user_entity_type"]
+ assert shared_state["user_entity_type_name"]
+ user_entity_type = shared_state["user_entity_type"]
+ user_entity_type_name = shared_state["user_entity_type_name"]
+
+ aiplatform.init(
+ project=e2e_base._PROJECT, location=e2e_base._LOCATION,
+ )
+
+ list_user_features = user_entity_type.list_features()
+ assert len(list_user_features) == 0
 
  # User Features
- user_age_feature_id = _TEST_USER_AGE_FEATURE_ID
  user_age_feature = user_entity_type.create_feature(
- feature_id=user_age_feature_id, value_type="INT64"
+ feature_id=_TEST_USER_AGE_FEATURE_ID, value_type="INT64"
  )
 
  get_user_age_feature = user_entity_type.get_feature(
- feature_id=user_age_feature_id
+ feature_id=_TEST_USER_AGE_FEATURE_ID
  )
  assert user_age_feature.resource_name == get_user_age_feature.resource_name
 
- user_gender_feature_id = _TEST_USER_GENDER_FEATURE_ID
  user_gender_feature = aiplatform.Feature.create(
- feature_id=user_gender_feature_id,
+ feature_id=_TEST_USER_GENDER_FEATURE_ID,
  value_type="STRING",
- entity_type_name=user_entity_type.resource_name,
+ entity_type_name=user_entity_type_name,
  )
 
  get_user_gender_feature = aiplatform.Feature(
@@ -127,66 +148,97 @@ def test_end_to_end(self, shared_state, caplog):
  user_gender_feature.resource_name == get_user_gender_feature.resource_name
  )
 
- user_liked_genres_feature_id = _TEST_USER_LIKED_GENRES_FEATURE_ID
+ list_user_features = aiplatform.Feature.list(
+ entity_type_name=user_entity_type_name
+ )
+ assert len(list_user_features) == 2
+
  user_feature_configs = {
- user_liked_genres_feature_id: {"value_type": "STRING_ARRAY"},
+ _TEST_USER_LIKED_GENRES_FEATURE_ID: {"value_type": "STRING_ARRAY"},
  }
  user_entity_type.batch_create_features(feature_configs=user_feature_configs)
 
  list_user_features = user_entity_type.list_features()
  assert len(list_user_features) == 3
 
- user_entity_type = user_entity_type.ingest_from_gcs(
+ def test_ingest_feature_values(self, shared_state, caplog):
+
+ assert shared_state["user_entity_type"]
+ user_entity_type = shared_state["user_entity_type"]
+
+ caplog.set_level(logging.INFO)
+
+ aiplatform.init(
+ project=e2e_base._PROJECT, location=e2e_base._LOCATION,
+ )
+
+ user_entity_type.ingest_from_gcs(
  feature_ids=[
- user_age_feature_id,
- user_gender_feature_id,
- user_liked_genres_feature_id,
+ _TEST_USER_AGE_FEATURE_ID,
+ _TEST_USER_GENDER_FEATURE_ID,
+ _TEST_USER_LIKED_GENRES_FEATURE_ID,
  ],
  feature_time="update_time",
  gcs_source_uris=_TEST_USERS_ENTITY_TYPE_GCS_SRC,
  gcs_source_type="avro",
  entity_id_field="user_id",
  worker_count=2,
- sync=False,
  )
 
- movie_title_feature_id = _TEST_MOVIE_TITLE_FEATURE_ID
- movie_genres_feature_id = _TEST_MOVIE_GENRES_FEATURE_ID
- movie_average_rating_id = _TEST_MOVIE_AVERAGE_RATING_FEATURE_ID
+ assert "EntityType feature values imported." in caplog.text
+
+ caplog.clear()
+
+ def test_ingest_feature_values_with_config(self, shared_state, caplog):
+
+ assert shared_state["movie_entity_type"]
+ movie_entity_type = shared_state["movie_entity_type"]
+
+ caplog.set_level(logging.INFO)
+
+ aiplatform.init(
+ project=e2e_base._PROJECT, location=e2e_base._LOCATION,
+ )
 
  movie_feature_configs = {
- movie_title_feature_id: {"value_type": "STRING"},
- movie_genres_feature_id: {"value_type": "STRING"},
- movie_average_rating_id: {"value_type": "DOUBLE"},
+ _TEST_MOVIE_TITLE_FEATURE_ID: {"value_type": "STRING"},
+ _TEST_MOVIE_GENRES_FEATURE_ID: {"value_type": "STRING"},
+ _TEST_MOVIE_AVERAGE_RATING_FEATURE_ID: {"value_type": "DOUBLE"},
  }
 
  list_movie_features = movie_entity_type.list_features()
  assert len(list_movie_features) == 0
 
- movie_entity_type = movie_entity_type.ingest_from_gcs(
+ movie_entity_type.ingest_from_gcs(
  feature_ids=[
- movie_title_feature_id,
- movie_genres_feature_id,
- movie_average_rating_id,
+ _TEST_MOVIE_TITLE_FEATURE_ID,
+ _TEST_MOVIE_GENRES_FEATURE_ID,
+ _TEST_MOVIE_AVERAGE_RATING_FEATURE_ID,
  ],
  feature_time="update_time",
  gcs_source_uris=_TEST_MOVIES_ENTITY_TYPE_GCS_SRC,
  gcs_source_type="avro",
  entity_id_field="movie_id",
  worker_count=2,
  batch_create_feature_configs=movie_feature_configs,
- sync=False,
  )
 
- user_entity_type.wait()
- movie_entity_type.wait()
-
- assert "EntityType feature values imported." in caplog.text
-
  list_movie_features = movie_entity_type.list_features()
  assert len(list_movie_features) == 3
 
- list_searched_features = aiplatform.Feature.search()
- assert len(list_searched_features) - base_list_searched_features == 6
+ assert "EntityType feature values imported." in caplog.text
 
  caplog.clear()
+
+ def test_search_features(self, shared_state):
+
+ assert shared_state["base_list_searched_features"] is not None
+
+ aiplatform.init(
+ project=e2e_base._PROJECT, location=e2e_base._LOCATION,
+ )
+
+ list_searched_features = aiplatform.Feature.search()
+ assert (
+ len(list_searched_features) - shared_state["base_list_searched_features"]
+ ) == 3