This is an automated email from the ASF dual-hosted git repository.
kamilbregula pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/master by this push:
new 3cc5756 Add unit tests for mlengine_operator_utils (#9702)
3cc5756 is described below
commit 3cc5756d04c7d3fd14b71e9ec6f5f906d5a24212
Author: chipmyersjr <[email protected]>
AuthorDate: Fri Jul 10 15:26:00 2020 -0700
Add unit tests for mlengine_operator_utils (#9702)
---
.../cloud/utils/test_mlengine_operator_utils.py | 277 +++++++++++++++++++++
tests/test_project_structure.py | 1 -
2 files changed, 277 insertions(+), 1 deletion(-)
diff --git a/tests/providers/google/cloud/utils/test_mlengine_operator_utils.py
b/tests/providers/google/cloud/utils/test_mlengine_operator_utils.py
new file mode 100644
index 0000000..326c14e
--- /dev/null
+++ b/tests/providers/google/cloud/utils/test_mlengine_operator_utils.py
@@ -0,0 +1,277 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import base64
+import json
+import unittest
+from datetime import datetime
+
+import dill
+import mock
+
+from airflow.exceptions import AirflowException
+from airflow.models import DAG
+from airflow.operators.python import PythonOperator
+from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.cloud.operators.dataflow import
DataflowCreatePythonJobOperator
+from airflow.providers.google.cloud.utils.mlengine_operator_utils import
create_evaluate_ops
+
+TASK_PREFIX = "test-task-prefix"
+TASK_PREFIX_PREDICTION = TASK_PREFIX + "-prediction"
+TASK_PREFIX_SUMMARY = TASK_PREFIX + "-summary"
+TASK_PREFIX_VALIDATION = TASK_PREFIX + "-validation"
+DATA_FORMAT = "TEXT"
+INPUT_PATHS = ["gs://path/to/input/file.json",
+ "gs://path/to/input/file2.json",
+ "gs://path/to/input/file3.json"]
+PREDICTION_PATH = "gs://path/to/output/predictions.json"
+BATCH_PREDICTION_JOB_ID = "test-batch-prediction-job-id"
+PROJECT_ID = "test-project-id"
+REGION = "test-region"
+DATAFLOW_OPTIONS = {
+ "project": "my-gcp-project",
+ "zone": "us-central1-f",
+ "stagingLocation": "gs://bucket/tmp/dataflow/staging/"
+}
+MODEL_URI = "gs://path/to/model/model"
+MODEL_NAME = "test-model-name"
+VERSION_NAME = "test-version-name"
+DAG_DEFAULT_ARGS = {
+ "project_id": PROJECT_ID,
+ "region": REGION,
+ "model_name": MODEL_NAME,
+ "version_name": VERSION_NAME,
+ "dataflow_default_options": DATAFLOW_OPTIONS
+}
+TEST_DAG = DAG(dag_id="test-dag-id", start_date=datetime(2000, 1, 1),
default_args=DAG_DEFAULT_ARGS)
+
+
+def get_metric_fn_and_keys():
+ import math
+
+ def error_and_squared_error(inst):
+ label = float(inst['input_label'])
+ classes = float(inst['classes'])
+ err = abs(classes - label)
+ squared_err = math.pow(classes - label, 2)
+ return err, squared_err
+
+ return error_and_squared_error, ['err', 'mse']
+
+
+METRIC_FN, METRIC_KEYS = get_metric_fn_and_keys()
+METRIC_FN_ENCODED = base64.b64encode(dill.dumps(METRIC_FN,
recurse=True)).decode()
+METRIC_KEYS_EXPECTED = ','.join(METRIC_KEYS)
+
+
+def validate_err_and_count(summary):
+ if summary['err'] > 0.2:
+ raise ValueError('Too high err>0.2; summary=%s' % summary)
+ if summary['mse'] > 0.05:
+ raise ValueError('Too high mse>0.05; summary=%s' % summary)
+ if summary['count'] < 1000:
+ raise ValueError('Too few instances<1000; summary=%s' % summary)
+ return summary
+
+
+class TestMlengineOperatorUtils(unittest.TestCase):
+ @mock.patch.object(PythonOperator, "set_upstream")
+ @mock.patch.object(DataflowCreatePythonJobOperator, "set_upstream")
+ def test_create_evaluate_ops(self, mock_dataflow, mock_python):
+ result = create_evaluate_ops(task_prefix=TASK_PREFIX,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+
metric_fn_and_keys=get_metric_fn_and_keys(),
+ validate_fn=validate_err_and_count,
+
batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
+ project_id=PROJECT_ID,
+ region=REGION,
+ dataflow_options=DATAFLOW_OPTIONS,
+ model_uri=MODEL_URI
+ )
+
+ evaluate_prediction, evaluate_summary, evaluate_validation = result
+
+ mock_dataflow.assert_called_once_with(evaluate_prediction)
+ mock_python.assert_called_once_with(evaluate_summary)
+
+ self.assertEqual(TASK_PREFIX_PREDICTION, evaluate_prediction.task_id)
+ self.assertEqual(PROJECT_ID, evaluate_prediction._project_id)
+ self.assertEqual(BATCH_PREDICTION_JOB_ID, evaluate_prediction._job_id)
+ self.assertEqual(REGION, evaluate_prediction._region)
+ self.assertEqual(DATA_FORMAT, evaluate_prediction._data_format)
+ self.assertEqual(INPUT_PATHS, evaluate_prediction._input_paths)
+ self.assertEqual(PREDICTION_PATH, evaluate_prediction._output_path)
+ self.assertEqual(MODEL_URI, evaluate_prediction._uri)
+
+ self.assertEqual(TASK_PREFIX_SUMMARY, evaluate_summary.task_id)
+ self.assertEqual(DATAFLOW_OPTIONS,
evaluate_summary.dataflow_default_options)
+ self.assertEqual(PREDICTION_PATH,
evaluate_summary.options["prediction_path"])
+ self.assertEqual(METRIC_FN_ENCODED,
evaluate_summary.options["metric_fn_encoded"])
+ self.assertEqual(METRIC_KEYS_EXPECTED,
evaluate_summary.options["metric_keys"])
+
+ self.assertEqual(TASK_PREFIX_VALIDATION, evaluate_validation.task_id)
+ self.assertEqual(PREDICTION_PATH,
evaluate_validation.templates_dict["prediction_path"])
+
+ @mock.patch.object(PythonOperator, "set_upstream")
+ @mock.patch.object(DataflowCreatePythonJobOperator, "set_upstream")
+ def test_create_evaluate_ops_model_and_version_name(self, mock_dataflow,
mock_python):
+ result = create_evaluate_ops(task_prefix=TASK_PREFIX,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+
metric_fn_and_keys=get_metric_fn_and_keys(),
+ validate_fn=validate_err_and_count,
+
batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
+ project_id=PROJECT_ID,
+ region=REGION,
+ dataflow_options=DATAFLOW_OPTIONS,
+ model_name=MODEL_NAME,
+ version_name=VERSION_NAME
+ )
+
+ evaluate_prediction, evaluate_summary, evaluate_validation = result
+
+ mock_dataflow.assert_called_once_with(evaluate_prediction)
+ mock_python.assert_called_once_with(evaluate_summary)
+
+ self.assertEqual(TASK_PREFIX_PREDICTION, evaluate_prediction.task_id)
+ self.assertEqual(PROJECT_ID, evaluate_prediction._project_id)
+ self.assertEqual(BATCH_PREDICTION_JOB_ID, evaluate_prediction._job_id)
+ self.assertEqual(REGION, evaluate_prediction._region)
+ self.assertEqual(DATA_FORMAT, evaluate_prediction._data_format)
+ self.assertEqual(INPUT_PATHS, evaluate_prediction._input_paths)
+ self.assertEqual(PREDICTION_PATH, evaluate_prediction._output_path)
+ self.assertEqual(MODEL_NAME, evaluate_prediction._model_name)
+ self.assertEqual(VERSION_NAME, evaluate_prediction._version_name)
+
+ self.assertEqual(TASK_PREFIX_SUMMARY, evaluate_summary.task_id)
+ self.assertEqual(DATAFLOW_OPTIONS,
evaluate_summary.dataflow_default_options)
+ self.assertEqual(PREDICTION_PATH,
evaluate_summary.options["prediction_path"])
+ self.assertEqual(METRIC_FN_ENCODED,
evaluate_summary.options["metric_fn_encoded"])
+ self.assertEqual(METRIC_KEYS_EXPECTED,
evaluate_summary.options["metric_keys"])
+
+ self.assertEqual(TASK_PREFIX_VALIDATION, evaluate_validation.task_id)
+ self.assertEqual(PREDICTION_PATH,
evaluate_validation.templates_dict["prediction_path"])
+
+ @mock.patch.object(PythonOperator, "set_upstream")
+ @mock.patch.object(DataflowCreatePythonJobOperator, "set_upstream")
+ def test_create_evaluate_ops_dag(self, mock_dataflow, mock_python):
+ result = create_evaluate_ops(task_prefix=TASK_PREFIX,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+
metric_fn_and_keys=get_metric_fn_and_keys(),
+ validate_fn=validate_err_and_count,
+
batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
+ dag=TEST_DAG
+ )
+
+ evaluate_prediction, evaluate_summary, evaluate_validation = result
+
+ mock_dataflow.assert_called_once_with(evaluate_prediction)
+ mock_python.assert_called_once_with(evaluate_summary)
+
+ self.assertEqual(TASK_PREFIX_PREDICTION, evaluate_prediction.task_id)
+ self.assertEqual(PROJECT_ID, evaluate_prediction._project_id)
+ self.assertEqual(BATCH_PREDICTION_JOB_ID, evaluate_prediction._job_id)
+ self.assertEqual(REGION, evaluate_prediction._region)
+ self.assertEqual(DATA_FORMAT, evaluate_prediction._data_format)
+ self.assertEqual(INPUT_PATHS, evaluate_prediction._input_paths)
+ self.assertEqual(PREDICTION_PATH, evaluate_prediction._output_path)
+ self.assertEqual(MODEL_NAME, evaluate_prediction._model_name)
+ self.assertEqual(VERSION_NAME, evaluate_prediction._version_name)
+
+ self.assertEqual(TASK_PREFIX_SUMMARY, evaluate_summary.task_id)
+ self.assertEqual(DATAFLOW_OPTIONS,
evaluate_summary.dataflow_default_options)
+ self.assertEqual(PREDICTION_PATH,
evaluate_summary.options["prediction_path"])
+ self.assertEqual(METRIC_FN_ENCODED,
evaluate_summary.options["metric_fn_encoded"])
+ self.assertEqual(METRIC_KEYS_EXPECTED,
evaluate_summary.options["metric_keys"])
+
+ self.assertEqual(TASK_PREFIX_VALIDATION, evaluate_validation.task_id)
+ self.assertEqual(PREDICTION_PATH,
evaluate_validation.templates_dict["prediction_path"])
+
+ @mock.patch.object(GCSHook, "download")
+ @mock.patch.object(PythonOperator, "set_upstream")
+ @mock.patch.object(DataflowCreatePythonJobOperator, "set_upstream")
+ def test_apply_validate_fn(self, mock_dataflow, mock_python,
mock_download):
+ result = create_evaluate_ops(task_prefix=TASK_PREFIX,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+
metric_fn_and_keys=get_metric_fn_and_keys(),
+ validate_fn=validate_err_and_count,
+
batch_prediction_job_id=BATCH_PREDICTION_JOB_ID,
+ project_id=PROJECT_ID,
+ region=REGION,
+ dataflow_options=DATAFLOW_OPTIONS,
+ model_uri=MODEL_URI
+ )
+
+ _, _, evaluate_validation = result
+
+ mock_download.return_value = json.dumps({
+ "err": 0.3,
+ "mse": 0.04,
+ "count": 1100
+ })
+ templates_dict = {"prediction_path": PREDICTION_PATH}
+ with self.assertRaises(ValueError) as context:
+ evaluate_validation.python_callable(templates_dict=templates_dict)
+
+ self.assertEqual("Too high err>0.2; summary={'err': 0.3, 'mse': 0.04,
'count': 1100}",
+ str(context.exception))
+ mock_download.assert_called_once_with("path",
"to/output/predictions.json/prediction.summary.json")
+
+ invalid_prediction_paths = ["://path/to/output/predictions.json",
"gs://", ""]
+
+ for path in invalid_prediction_paths:
+ templates_dict = {"prediction_path": path}
+ with self.assertRaises(ValueError) as context:
+
evaluate_validation.python_callable(templates_dict=templates_dict)
+ self.assertEqual("Wrong format prediction_path:",
str(context.exception)[:29])
+
+ def test_invalid_task_prefix(self):
+ invalid_task_prefix_values = ["test-task-prefix&",
"~test-task-prefix", "test-task(-prefix"]
+
+ for invalid_task_prefix_value in invalid_task_prefix_values:
+ with self.assertRaises(AirflowException):
+ create_evaluate_ops(task_prefix=invalid_task_prefix_value,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+
metric_fn_and_keys=get_metric_fn_and_keys(),
+ validate_fn=validate_err_and_count)
+
+ def test_non_callable_metric_fn(self):
+ with self.assertRaises(AirflowException):
+ create_evaluate_ops(task_prefix=TASK_PREFIX,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+ metric_fn_and_keys=("error_and_squared_error",
['err', 'mse']),
+ validate_fn=validate_err_and_count)
+
+ def test_non_callable_validate_fn(self):
+ with self.assertRaises(AirflowException):
+ create_evaluate_ops(task_prefix=TASK_PREFIX,
+ data_format=DATA_FORMAT,
+ input_paths=INPUT_PATHS,
+ prediction_path=PREDICTION_PATH,
+ metric_fn_and_keys=get_metric_fn_and_keys(),
+ validate_fn="validate_err_and_count")
diff --git a/tests/test_project_structure.py b/tests/test_project_structure.py
index 0166868..bd5941e 100644
--- a/tests/test_project_structure.py
+++ b/tests/test_project_structure.py
@@ -36,7 +36,6 @@ MISSING_TEST_FILES = {
'tests/providers/google/cloud/transfers/test_sql_to_gcs.py',
'tests/providers/google/cloud/utils/test_field_sanitizer.py',
'tests/providers/google/cloud/utils/test_field_validator.py',
- 'tests/providers/google/cloud/utils/test_mlengine_operator_utils.py',
'tests/providers/google/cloud/utils/test_mlengine_prediction_summary.py',
'tests/providers/jenkins/hooks/test_jenkins.py',
'tests/providers/microsoft/azure/sensors/test_azure_cosmos.py',
