server: tests: minor fix missing param.

phymbert · Feb 22, 2024 · cba6d4e · cba6d4e
1 parent 51f5274
commit cba6d4e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/server/tests/features/steps/steps.py b/examples/server/tests/features/steps/steps.py
@@ -282,7 +282,7 @@ async def step_all_prompts_are_predicted(context, n_predict):
     await all_prompts_are_predicted(context, expected_predicted_n)
 
 
-async def all_prompts_are_predicted(context, expected_predicted_n):
+async def all_prompts_are_predicted(context, expected_predicted_n=None):
     n_completions = await gather_concurrent_completions_tasks(context)
     assert n_completions > 0
     for i in range(n_completions):
@@ -567,7 +567,8 @@ async def wait_for_health_status(context,
                 # Sometimes health requests are triggered after completions are predicted
                 if expected_http_status_code == 503:
                     if len(context.completions) == 0:
-                        print("\x1b[5;37;43mWARNING: forcing concurrent completions task, busy health check missed")
+                        print("\x1b[5;37;43mWARNING: forcing concurrents completions tasks,"
+                              " busy health check missed\x1b[0m")
                         n_completions = await gather_concurrent_completions_tasks(context)
                         if n_completions > 0:
                             return