diff --git a/translate/automl/model_test.py b/translate/automl/model_test.py index 0d37a85c674..eca489dbd9a 100644 --- a/translate/automl/model_test.py +++ b/translate/automl/model_test.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import datetime import os +import uuid from google.cloud import automl_v1beta1 as automl import pytest @@ -26,19 +26,25 @@ compute_region = "us-central1" -@pytest.mark.skip(reason="creates too many models") -def test_model_create_status_delete(capsys): - # create model +@pytest.fixture +def model(): client = automl.AutoMlClient() - model_name = "test_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S") + model_name = "test_{}".format(uuid.uuid4()).replace("-", "")[:32] project_location = client.location_path(project_id, compute_region) my_model = { "display_name": model_name, "dataset_id": "3876092572857648864", "translation_model_metadata": {"base_model": ""}, } - response = client.create_model(project_location, my_model) - operation_name = response.operation.name + operation = client.create_model(project_location, my_model) + + yield operation + + operation.cancel() + + +def test_model_create_status_delete(capsys, model): + operation_name = model.operation.name assert operation_name # get operation status @@ -46,9 +52,6 @@ def test_model_create_status_delete(capsys): out, _ = capsys.readouterr() assert "Operation status: " in out - # cancel operation - response.cancel() - def test_model_list_get_evaluate(capsys): # list models diff --git a/vision/automl/model_test.py b/vision/automl/model_test.py index a272ce92334..bae516664c0 100644 --- a/vision/automl/model_test.py +++ b/vision/automl/model_test.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import datetime import os +import uuid from google.cloud import automl_v1beta1 as automl import pytest @@ -24,11 +24,10 @@ compute_region = "us-central1" -@pytest.mark.skip(reason="creates too many models") -def test_model_create_status_delete(capsys): - # create model +@pytest.fixture +def model(): client = automl.AutoMlClient() - model_name = "test_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S") + model_name = "test_{}".format(uuid.uuid4()).replace("-", "")[:32] project_location = client.location_path(project_id, compute_region) my_model = { "display_name": model_name, @@ -36,8 +35,12 @@ def test_model_create_status_delete(capsys): "image_classification_model_metadata": {"train_budget": 24}, } response = client.create_model(project_location, my_model) - operation_name = response.operation.name - assert operation_name - # cancel operation + yield response + response.cancel() + + +def test_model_create_status_delete(capsys, model): + operation_name = model.operation.name + assert operation_name