diff --git a/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml index 31d29ef04b..4d02bd8d14 100644 --- a/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml +++ b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml @@ -1,5 +1,5 @@ build: python_version: "3.11" python_packages: - - "pillow==8.3.2" + - "tqdm==4.66.6" predict: "predict.py:Predictor" diff --git a/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py index 5ce00f182c..4eb76c5eba 100644 --- a/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py +++ b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py @@ -1,22 +1,12 @@ -from cog import BasePredictor, File -import logging -import os -from PIL import Image +from cog import BasePredictor +import tqdm +import time class Predictor(BasePredictor): - def setup(self): - logFormatter = logging.Formatter("%(asctime)s [%(threadName)-12.12s] [%(levelname)-5.5s] %(message)s") - rootLogger = logging.getLogger() - fileHandler = logging.FileHandler("{0}/{1}.log".format(os.path.dirname(__file__), "mylog.log")) - fileHandler.setFormatter(logFormatter) - rootLogger.addHandler(fileHandler) - - consoleHandler = logging.StreamHandler() - consoleHandler.setFormatter(logFormatter) - rootLogger.addHandler(consoleHandler) - - def predict(self) -> list[Image]: - logging.info("Do some logging.") - return [] + def predict(self) -> str: + print("Starting Predict:") + for i in tqdm.tqdm(range(10), total=10): + time.sleep(i) + return "Hello World" diff --git a/test-integration/test_integration/test_predict.py b/test-integration/test_integration/test_predict.py index e417d56537..a25f4e2aa6 100644 --- a/test-integration/test_integration/test_predict.py +++ b/test-integration/test_integration/test_predict.py @@ -334,7 +334,6 @@ def test_file_logging_stream_drain(tmpdir_factory): cmd, cwd=out_dir, check=True, - capture_output=True, text=True, timeout=DEFAULT_TIMEOUT, )