From b8263dd6084954b87d0521d239a9cbf44270cbec Mon Sep 17 00:00:00 2001
From: Raphael Glon <oOraph@users.noreply.github.com>
Date: Thu, 30 Nov 2023 10:56:43 +0100
Subject: [PATCH] fix diffusers test

Signed-off-by: Raphael Glon <oOraph@users.noreply.github.com>
---
 tests/test_dockers.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/test_dockers.py b/tests/test_dockers.py
index 0df16d3f..7280f635 100644
--- a/tests/test_dockers.py
+++ b/tests/test_dockers.py
@@ -313,12 +313,16 @@ def test_diffusers(self):
             "diffusers",
             "text-to-image",
             "hf-internal-testing/tiny-stable-diffusion-pipe",
+            {"num_inference_steps": 1},
+            custom_environment={"UVICORN_TIMEOUT": "1200"},
         )
         self.framework_docker_test(
             "diffusers",
             "image-to-image",
             "hf-internal-testing/tiny-controlnet",
+            {"num_inference_steps": 1},
             timeout=600,
+            custom_environment={"UVICORN_TIMEOUT": "1200"},
         )
         self.framework_invalid_test("diffusers")
 
@@ -479,6 +483,7 @@ def framework_docker_test(
             Any
         ] = None,  # if given, check inference with this specific input
         timeout=60,
+        custom_environment: Optional[dict] = None,
     ):
         tag = self.create_docker(framework)
         run_docker_command = [
@@ -490,6 +495,12 @@ def framework_docker_test(
             f"TASK={task}",
             "-e",
             f"MODEL_ID={model_id}",
+        ]
+
+        if custom_environment:
+            for k, v in custom_environment.items():
+                run_docker_command += ["-e", f"{k}={v}"]
+        run_docker_command += [
             "-v",
             "/tmp:/data",
             "-t",