From 1aa5ad91505f988f7626b8d4e631733f6580ab95 Mon Sep 17 00:00:00 2001 From: Pierrick HYMBERT Date: Sat, 2 Mar 2024 19:30:19 +0100 Subject: [PATCH] server: tests: fix re content --- examples/server/tests/features/steps/steps.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/examples/server/tests/features/steps/steps.py b/examples/server/tests/features/steps/steps.py index 5642122f9..403b0eae7 100644 --- a/examples/server/tests/features/steps/steps.py +++ b/examples/server/tests/features/steps/steps.py @@ -812,13 +812,21 @@ def assert_n_tokens_predicted(completion_response, expected_predicted_n=None, re n_predicted = completion_response['timings']['predicted_n'] assert len(content) > 0, "no token predicted" if re_content is not None: - re_content = f'^(.*)({re_content})(.*)$' p = re.compile(re_content, flags=RegexFlag.IGNORECASE | RegexFlag.MULTILINE | RegexFlag.DOTALL) - match = p.match(content) - assert match and len(match.groups()) >= 3, f'/{re_content}/g must match ```{content}```' + matches = p.finditer(content) + last_match = 0 + highlighted = '' + for match in matches: + start, end = match.span() + highlighted += content[last_match: start] + highlighted += '\x1b[33m' + highlighted += content[start: end] + highlighted += '\x1b[0m' + last_match = end + highlighted += content[last_match:] if 'DEBUG' in os.environ and os.environ['DEBUG'] == 'ON': - highlighted = p.sub(r"\1\2\3", content).replace('', '\x1b[33m').replace('', '\x1b[0m') - print(f"Checking completion response: {highlighted}\n") + print(f"Checking completion response: {highlighted}\n") + assert last_match > 0, f'/{re_content}/ must match ```{highlighted}```' if expected_predicted_n and expected_predicted_n > 0: assert n_predicted == expected_predicted_n, (f'invalid number of tokens predicted:' f' {n_predicted} <> {expected_predicted_n}')