From 60d61746c75330eb36df29078968b642d6543956 Mon Sep 17 00:00:00 2001
From: Aaron Davidson <aaron@databricks.com>
Date: Wed, 15 Apr 2020 16:23:34 -0700
Subject: [PATCH] pytest --color=yes

---
 travis/run-large-python-tests.sh | 42 ++++++++++++++++----------------
 travis/run-small-python-tests.sh |  2 +-
 travis/stage-python3.sh          |  2 +-
 3 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/travis/run-large-python-tests.sh b/travis/run-large-python-tests.sh
index 31287abe4cd4be..16df5cab554919 100755
--- a/travis/run-large-python-tests.sh
+++ b/travis/run-large-python-tests.sh
@@ -16,7 +16,7 @@ else
 fi
 
 # NB: Also add --ignore'd tests to run-small-python-tests.sh
-pytest tests --large --ignore=tests/examples --ignore=tests/h2o --ignore=tests/keras \
+pytest tests --color=yes --large --ignore=tests/examples --ignore=tests/h2o --ignore=tests/keras \
   --ignore=tests/pytorch --ignore=tests/pyfunc --ignore=tests/sagemaker --ignore=tests/sklearn \
   --ignore=tests/spark --ignore=tests/tensorflow --ignore=tests/azureml --ignore=tests/onnx \
   --ignore=tests/keras_autolog --ignore=tests/tensorflow_autolog --ignore=tests/gluon \
@@ -24,29 +24,29 @@ pytest tests --large --ignore=tests/examples --ignore=tests/h2o --ignore=tests/k
   --ignore tests/spark_autologging --ignore=tests/models
 # Run ML framework tests in their own Python processes to avoid OOM issues due to per-framework
 # overhead
-pytest --verbose tests/pytorch --large
-pytest --verbose tests/h2o --large
-pytest --verbose tests/onnx --large
-pytest --verbose tests/pyfunc --large
-pytest --verbose tests/sagemaker --large
-pytest --verbose tests/sagemaker/mock --large
-pytest --verbose tests/sklearn --large
-pytest --verbose tests/spark --large
-pytest --verbose tests/tensorflow/test_tensorflow_model_export.py --large
-pytest --verbose tests/tensorflow_autolog/test_tensorflow_autolog.py --large
-pytest --verbose tests/azureml --large
-pytest --verbose tests/models --large
-pytest --verbose tests/xgboost --large
-pytest --verbose tests/lightgbm --large
+pytest --color=yes --verbose tests/pytorch --large
+pytest --color=yes --verbose tests/h2o --large
+pytest --color=yes --verbose tests/onnx --large
+pytest --color=yes --verbose tests/pyfunc --large
+pytest --color=yes --verbose tests/sagemaker --large
+pytest --color=yes --verbose tests/sagemaker/mock --large
+pytest --color=yes --verbose tests/sklearn --large
+pytest --color=yes --verbose tests/spark --large
+pytest --color=yes --verbose tests/tensorflow/test_tensorflow_model_export.py --large
+pytest --color=yes --verbose tests/tensorflow_autolog/test_tensorflow_autolog.py --large
+pytest --color=yes --verbose tests/azureml --large
+pytest --color=yes --verbose tests/models --large
+pytest --color=yes --verbose tests/xgboost --large
+pytest --color=yes --verbose tests/lightgbm --large
 # TODO(smurching) Unpin TensorFlow dependency version once test failures with TF 2.1.0 have been
 # fixed
 pip install 'tensorflow==2.0.0'
-pytest --verbose tests/tensorflow/test_tensorflow2_model_export.py --large
-pytest --verbose tests/tensorflow_autolog/test_tensorflow2_autolog.py --large
-pytest --verbose tests/keras --large
-pytest --verbose tests/keras_autolog --large
-pytest --verbose tests/gluon --large
-pytest --verbose tests/gluon_autolog --large
+pytest --color=yes --verbose tests/tensorflow/test_tensorflow2_model_export.py --large
+pytest --color=yes --verbose tests/tensorflow_autolog/test_tensorflow2_autolog.py --large
+pytest --color=yes --verbose tests/keras --large
+pytest --color=yes --verbose tests/keras_autolog --large
+pytest --color=yes --verbose tests/gluon --large
+pytest --color=yes --verbose tests/gluon_autolog --large
 
 # Run Spark autologging tests
 ./travis/test-spark-autologging.sh
diff --git a/travis/run-small-python-tests.sh b/travis/run-small-python-tests.sh
index 4ed60ccbf0cc26..a39daaadfcff94 100755
--- a/travis/run-small-python-tests.sh
+++ b/travis/run-small-python-tests.sh
@@ -7,7 +7,7 @@ trap 'err=1' ERR
 export MLFLOW_HOME=$(pwd)
 
 # NB: Also add --ignore'd tests to run-large-python-tests.sh
-pytest --cov=mlflow --verbose --ignore=tests/h2o --ignore=tests/keras \
+pytest --color=yes --cov=mlflow --verbose --ignore=tests/h2o --ignore=tests/keras \
   --ignore=tests/pytorch --ignore=tests/pyfunc --ignore=tests/sagemaker --ignore=tests/sklearn \
   --ignore=tests/spark --ignore=tests/tensorflow --ignore=tests/keras_autolog \
   --ignore=tests/tensorflow_autolog --ignore tests/azureml --ignore tests/onnx \
diff --git a/travis/stage-python3.sh b/travis/stage-python3.sh
index 1b055bab0b7864..fccd04f99b5417 100755
--- a/travis/stage-python3.sh
+++ b/travis/stage-python3.sh
@@ -17,7 +17,7 @@ fi
 CHANGED_FILES=$(git diff --name-only master..HEAD | grep "tests/examples\|examples") || true
 if [[ "$TRAVIS_EVENT_TYPE" == "cron" || "$CHANGED_FILES" == *"examples"* ]] && [[ "$TRAVIS_BUILD_STAGE_NAME" == "Nightly" ]]
 then
-  pytest --verbose tests/examples --large;
+  pytest --color=yes --verbose tests/examples --large;
 fi
 if [[ "$TRAVIS_EVENT_TYPE" == "cron" || "$CHANGED_FILES" == *"Dockerfile"* ]] && [[ "$TRAVIS_BUILD_STAGE_NAME" == "Nightly" ]]
 then