From 6abd658a019ca636e9b60137874e130af1462959 Mon Sep 17 00:00:00 2001
From: Michele Pangrazzi <xmikex83@gmail.com>
Date: Thu, 28 Nov 2024 10:12:32 +0100
Subject: [PATCH] Use ChatMessage class methods to build messages

---
 .../builders/test_answer_builder.py           | 32 +++++++------------
 1 file changed, 11 insertions(+), 21 deletions(-)

diff --git a/test/components/builders/test_answer_builder.py b/test/components/builders/test_answer_builder.py
index 9c220ef477..3b53c7251c 100644
--- a/test/components/builders/test_answer_builder.py
+++ b/test/components/builders/test_answer_builder.py
@@ -165,10 +165,8 @@ def test_run_with_reference_pattern_set_at_runtime(self):
     def test_run_with_chat_message_replies_without_pattern(self):
         component = AnswerBuilder()
         replies = [
-            ChatMessage(
-                content="Answer: AnswerString",
-                role=ChatRole.ASSISTANT,
-                name=None,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -194,10 +192,8 @@ def test_run_with_chat_message_replies_without_pattern(self):
     def test_run_with_chat_message_replies_with_pattern(self):
         component = AnswerBuilder(pattern=r"Answer: (.*)")
         replies = [
-            ChatMessage(
-                content="Answer: AnswerString",
-                role=ChatRole.ASSISTANT,
-                name=None,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -223,10 +219,8 @@ def test_run_with_chat_message_replies_with_pattern(self):
     def test_run_with_chat_message_replies_with_documents(self):
         component = AnswerBuilder(reference_pattern="\\[(\\d+)\\]")
         replies = [
-            ChatMessage(
-                content="Answer: AnswerString[2]",
-                role=ChatRole.ASSISTANT,
-                name=None,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString[2]",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -257,10 +251,8 @@ def test_run_with_chat_message_replies_with_documents(self):
     def test_run_with_chat_message_replies_with_pattern_set_at_runtime(self):
         component = AnswerBuilder(pattern="unused pattern")
         replies = [
-            ChatMessage(
-                content="Answer: AnswerString",
-                role=ChatRole.ASSISTANT,
-                name=None,
+            ChatMessage.from_assistant(
+                "Answer: AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -286,10 +278,8 @@ def test_run_with_chat_message_replies_with_pattern_set_at_runtime(self):
     def test_run_with_chat_message_replies_with_meta_set_at_run_time(self):
         component = AnswerBuilder()
         replies = [
-            ChatMessage(
-                content="AnswerString",
-                role=ChatRole.ASSISTANT,
-                name=None,
+            ChatMessage.from_assistant(
+                "AnswerString",
                 meta={
                     "model": "gpt-4o-mini",
                     "index": 0,
@@ -315,7 +305,7 @@ def test_run_with_chat_message_replies_with_meta_set_at_run_time(self):
 
     def test_run_with_chat_message_no_meta_with_meta_set_at_run_time(self):
         component = AnswerBuilder()
-        replies = [ChatMessage(content="AnswerString", role=ChatRole.ASSISTANT, name=None, meta={})]
+        replies = [ChatMessage.from_assistant("AnswerString", meta={})]
         output = component.run(query="test query", replies=replies, meta=[{"test": "meta"}])
         answers = output["answers"]
         assert len(answers) == 1