Fixed review comment

Author: ryota-cloud

metadata-ingestion/src/datahub/ingestion/source/vertexai.py

@@ -273,7 +273,7 @@ def _gen_training_job_mcps(
             if job.create_time
             else datetime_to_ts_millis(datetime.now())
         )
-        created_actor = f"urn:li:platformResource:{self.platform}"
+        created_actor = "urn:li:corpuser:datahub"
 
         aspects: List[_Aspect] = list()
         aspects.append(
@@ -402,7 +402,7 @@ def _search_dataset(self, dataset_id: str) -> Optional[VertexAiResourceNoun]:
                 for ds in dataset_class.list():
                     self.datasets[ds.name] = ds
 
-        return self.datasets[dataset_id] if dataset_id in self.datasets else None
+        return self.datasets.get(dataset_id)
 
     def _get_input_dataset_mcps(
         self, job_meta: TrainingJobMetadata

metadata-ingestion/tests/integration/vertexai/test_vertexai.py

@@ -3,7 +3,6 @@
 from typing import Any, Dict, TypeVar
 from unittest.mock import patch
 
-import pytest
 from pytest import Config
 
 from datahub.ingestion.run.pipeline import Pipeline
@@ -22,11 +21,6 @@
 REGION = "us-west2"
 
 
-@pytest.fixture
-def sink_file_path(tmp_path: Path) -> str:
-    return str(tmp_path / "vertexai_source_mcps.json")
-
-
 def get_pipeline_config(sink_file_path: str) -> Dict[str, Any]:
     source_type = "vertexai"
     return {
@@ -47,7 +41,7 @@ def get_pipeline_config(sink_file_path: str) -> Dict[str, Any]:
     }
 
 
-def test_vertexai_source_ingestion(pytestconfig: Config, sink_file_path: str) -> None:
+def test_vertexai_source_ingestion(pytestconfig: Config, tmp_path: Path) -> None:
     with contextlib.ExitStack() as exit_stack:
         # Mock the Vertex API with empty list
         for func_to_mock in [
@@ -99,6 +93,7 @@ def test_vertexai_source_ingestion(pytestconfig: Config, sink_file_path: str) ->
             / "tests/integration/vertexai/vertexai_mcps_golden.json"
         )
 
+        sink_file_path = str(tmp_path / "vertexai_source_mcps.json")
         print(f"Output mcps file path: {str(sink_file_path)}")
         print(f"Golden file path: {str(golden_file_path)}")