fix: Fix message to Json parse issue in Dataset class

vertex-sdk-bot · copybara-github · commit 066f32df4e1c · 2024-01-18T17:46:39.000-08:00
PiperOrigin-RevId: 599675216
diff --git a/google/cloud/aiplatform/datasets/dataset.py b/google/cloud/aiplatform/datasets/dataset.py
@@ -805,7 +805,7 @@ def export_data_for_custom_training(
                 saved_query_id,
                 annotation_schema_uri,
                 split,
-            )
+            )._pb
         )
 
     def update(
diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py
@@ -405,9 +405,14 @@ def test_export_data_for_custom_training(self, staging_bucket):
         )
 
         # Ensure three output paths (training, validation and test) are provided
-        assert len(export_data_response["exported_files"]) == 3
+        assert len(export_data_response["exportedFiles"]) == 3
         # Ensure data stats are calculated and correct
-        assert export_data_response["data_stats"]["training_data_items_count"] == 40
+        assert int(export_data_response["dataStats"]["trainingDataItemsCount"]) == 40
+        assert int(export_data_response["dataStats"]["validationDataItemsCount"]) == 5
+        assert int(export_data_response["dataStats"]["testDataItemsCount"]) == 5
+        assert int(export_data_response["dataStats"]["trainingAnnotationsCount"]) == 40
+        assert int(export_data_response["dataStats"]["validationAnnotationsCount"]) == 5
+        assert int(export_data_response["dataStats"]["testAnnotationsCount"]) == 5
 
     def test_update_dataset(self):
         """Create a new dataset and use update() method to change its display_name, labels, and description.

Original file line number	Diff line number	Diff line change
`@@ -805,7 +805,7 @@ def export_data_for_custom_training(`
`805`	`805`	`saved_query_id,`
`806`	`806`	`annotation_schema_uri,`
`807`	`807`	`split,`
`808`		`- )`
	`808`	`+ )._pb`
`809`	`809`	`)`
`810`	`810`
`811`	`811`	`def update(`