Skip to content

Commit

Permalink
delete test case using real model
Browse files Browse the repository at this point in the history
  • Loading branch information
ryota-cloud committed Mar 4, 2025
1 parent 5765025 commit 4b09365
Showing 1 changed file with 0 additions and 63 deletions.
63 changes: 0 additions & 63 deletions metadata-ingestion/tests/unit/test_vertexai_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,35 +126,6 @@ def source() -> VertexAISource:
config=VertexAIConfig(project_id=PROJECT_ID, region=REGION),
)


def gen_real_model(source: VertexAISource) -> Model:
"""
Fixture for the model that is actually registered in the Vertex AI Model Registry
Use mock_model for local testing purpose, but this fixture is provided to use real model for debugging.
Replace model name with your real model when using this fixture.
"""
model_name = "projects/872197881936/locations/us-west2/models/3583871344875405312"
return Model(model_name=model_name)


def gen_real_autoML_tabular_job(source: VertexAISource) -> _TrainingJob:
"""
Fixture for the training job that is actually registered in the Vertex AI Model Registry
Use mock_training_job for local testing purpose, but this fixture is provided to use real training job for debugging.
Replace training job name with your real training job when using this fixture.
"""

# Initialize the AI Platform client
aiplatform.init(project=source.config.project_id, location=source.config.region)

# Retrieve the custom training job by its resource name
# resource_name format 'projects/your-project-id/locations/your-location/trainingPipelines/your-training-job-id')
job = aiplatform.AutoMLTabularTrainingJob.get(
resource_name="projects/872197881936/locations/us-west2/trainingPipelines/5401695018589093888"
)
return job


@patch("google.cloud.aiplatform.Model.list")
def test_get_ml_model_workunits(mock_list: List[Model], source: VertexAISource) -> None:
mock_models = gen_mock_models()
Expand Down Expand Up @@ -419,37 +390,3 @@ def test_make_job_urn(source: VertexAISource) -> None:
source._make_job_urn(mock_training_job)
== f"{builder.make_data_process_instance_urn(source._make_vertexai_job_name(mock_training_job.name))}"
)


@pytest.mark.skip(reason="Skipping, this is for debugging purpose")
def test_real_model_workunit(source: VertexAISource) -> None:
"""
Disabled as default
Use real model registered in the Vertex AI Model Registry
"""
real_model = gen_real_model(source)
model_version = gen_mock_model_version(real_model)
for wu in source._gen_ml_model_workunits(
model=real_model, model_version=model_version
):
assert hasattr(wu.metadata, "aspect")
aspect = wu.metadata.aspect
assert isinstance(aspect, MLModelProperties)
# aspect is MLModelPropertiesClass
assert aspect.description == model_version.version_description
assert aspect.date == model_version.version_create_time
assert aspect.hyperParams is None
assert aspect.trainingMetrics is None


@pytest.mark.skip(reason="Skipping, this is for debugging purpose")
def test_real_get_data_process_properties(source: VertexAISource) -> None:
real_autoML_tabular_job = gen_real_autoML_tabular_job(source)
for wu in source._get_training_job_workunits(real_autoML_tabular_job):
assert hasattr(wu.metadata, "aspect")
aspect = wu.metadata.aspect
if isinstance(aspect, DataProcessInstancePropertiesClass):
# aspect is DataProcessInstancePropertiesClass
assert aspect.externalUrl == source._make_job_external_url(
real_autoML_tabular_job
)

0 comments on commit 4b09365

Please sign in to comment.