Test tqdm stream redirection

test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml

@@ -0,0 +1,5 @@
+build:
+  python_version: "3.11"
+  python_packages:
+    - "tqdm==4.66.6"
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py

@@ -0,0 +1,12 @@
+from cog import BasePredictor
+import tqdm
+import time
+
+
+class Predictor(BasePredictor):
+
+    def predict(self) -> str:
+        print("Starting Predict:")
+        for i in tqdm.tqdm(range(10), total=10):
+            time.sleep(i)
+        return "Hello World"

test-integration/test_integration/test_predict.py

@@ -322,3 +322,18 @@ def test_predict_with_subprocess_in_setup(fixture_name):
             assert response.status_code == 200, str(response)
 
         assert busy_count < 10
+
+
+def test_file_logging_stream_drain(tmpdir_factory):
+    project_dir = Path(__file__).parent / "fixtures/file-logging-stream-drain-project"
+    out_dir = pathlib.Path(tmpdir_factory.mktemp("project"))
+    shutil.copytree(project_dir, out_dir, dirs_exist_ok=True)
+    cmd = ["cog", "--debug", "predict"]
+
+    subprocess.run(
+        cmd,
+        cwd=out_dir,
+        check=True,
+        text=True,
+        timeout=DEFAULT_TIMEOUT,
+    )