# from lionagi.core.branch.util import MessageUtil
# from lionagi.core.messages.schema import System, Instruction, Response

# import unittest
# import pandas as pd
# import json
# from datetime import datetime


# class TestCreateMessage(unittest.TestCase):

# 	def test_create_system_message(self):
# 		"""Test creating a System message."""
# 		system_info = {"system_info": "System information"}
# 		message = MessageUtil.create_message(
# 			system=system_info["system_info"]
# 			)
# 		self.assertIsInstance(message, System)
# 		self.assertEqual(message.content, system_info)

# 	def test_create_instruction_message(self):
# 		"""Test creating an Instruction message with context."""
# 		instruction_info = {"task": "Do something"}
# 		context = {"additional": "context"}
# 		message = MessageUtil.create_message(
# 			instruction=instruction_info, context=context
# 		)
# 		self.assertIsInstance(message, Instruction)
# 		self.assertEqual(message.content["instruction"], instruction_info)
# 		self.assertEqual(message.content["context"], context)

# 	def test_create_response_message(self):
# 		"""Test creating a Response message."""
# 		response_info = {"message": {"content": "This is a response"}}
# 		message = MessageUtil.create_message(response=response_info)
# 		self.assertIsInstance(message, Response)
# 		self.assertEqual(
# 			message.content["response"], response_info["message"]["content"]
# 		)

# 	def test_error_on_multiple_roles(self):
# 		"""Test error is raised when multiple roles are provided."""
# 		with self.assertRaises(ValueError):
# 			MessageUtil.create_message(
# 				system={"info": "info"}, instruction={"task": "task"}
# 			)

# 	def test_return_existing_base_message_instance(self):
# 		"""Test returning an existing BaseMessage instance if provided."""
# 		existing_message = System(system={"info": "Already created"})
# 		message = MessageUtil.create_message(system=existing_message)
# 		self.assertEqual(message.content, existing_message.content)


# class TestValidateMessages(unittest.TestCase):

# 	# def test_validate_messages_correct_format(self):
# 	#     """Test messages DataFrame with the correct format."""
# 	#     messages = pd.DataFrame({
# 	#         "node_id": ["1"],
# 	#         "role": ["user"],
# 	#         "sender": ["test"],
# 	#         "timestamp": ["2020-01-01T00:00:00"],
# 	#         "content": ['{"message": "test"}']
# 	#     })
# 	#     self.assertTrue(MessageUtil.validate_messages(messages))

# 	def test_validate_messages_incorrect_columns(self):
# 		"""Test messages DataFrame with incorrect columns raises ValueError."""
# 		messages = pd.DataFrame(
# 			{
# 				"id": ["1"], "type": ["user"], "source": ["test"],
# 				"time": ["2020-01-01T00:00:00"],
# 				"data": ['{"message": "test"}'],
# 			}
# 		)
# 		with self.assertRaises(ValueError):
# 			MessageUtil.validate_messages(messages)

# 	def test_validate_messages_null_values(self):
# 		"""Test messages DataFrame with null values raises ValueError."""
# 		messages = pd.DataFrame(
# 			{
# 				"node_id": [None], "role": ["user"], "sender": ["test"],
# 				"timestamp": ["2020-01-01T00:00:00"],
# 				"content": ['{"message": "test"}'],
# 			}
# 		)
# 		with self.assertRaises(ValueError):
# 			MessageUtil.validate_messages(messages)


# class TestSignMessage(unittest.TestCase):

# 	def test_sign_message(self):
# 		"""Test signing message content with sender."""
# 		messages = pd.DataFrame(
# 			{
# 				"node_id": ["1"], "role": ["user"], "sender": ["test"],
# 				"timestamp": ["2020-01-01T00:00:00"],
# 				"content": ["Original message"],
# 			}
# 		)
# 		sender = "system"
# 		signed_messages = MessageUtil.sign_message(messages, sender)
# 		expected_content = "Sender system: Original message"
# 		self.assertEqual(signed_messages["content"][0], expected_content)

# 	def test_sign_message_invalid_sender(self):
# 		"""Test signing message with an invalid sender raises ValueError."""
# 		messages = pd.DataFrame(
# 			{
# 				"node_id": ["1"], "role": ["user"], "sender": ["test"],
# 				"timestamp": ["2020-01-01T00:00:00"],
# 				"content": ["Original message"],
# 			}
# 		)
# 		with self.assertRaises(ValueError):
# 			MessageUtil.sign_message(messages, None)


# class TestFilterMessagesBy(unittest.TestCase):

# 	def setUp(self):
# 		self.messages = pd.DataFrame(
# 			{
# 				"node_id": ["1", "2"], "role": ["user", "assistant"],
# 				"sender": ["test", "assistant"],
# 				"timestamp": [datetime(2020, 1, 1), datetime(2020, 1, 2)],
# 				"content": ['{"message": "test"}', '{"response": "ok"}'],
# 			}
# 		)

# 	def test_filter_by_role(self):
# 		"""Test filtering messages by role."""
# 		filtered = MessageUtil.filter_messages_by(
# 			self.messages, role="assistant"
# 			)
# 		self.assertEqual(len(filtered), 1)
# 		self.assertEqual(filtered.iloc[0]["sender"], "assistant")

# 	def test_filter_by_sender(self):
# 		"""Test filtering messages by sender."""
# 		filtered = MessageUtil.filter_messages_by(
# 			self.messages, sender="test"
# 			)
# 		self.assertEqual(len(filtered), 1)
# 		self.assertEqual(filtered.iloc[0]["sender"], "test")

# 	def test_filter_by_time_range(self):
# 		"""Test filtering messages by time range."""
# 		start_time = datetime(2020, 1, 1, 12)
# 		end_time = datetime(2020, 1, 2, 12)
# 		filtered = MessageUtil.filter_messages_by(
# 			self.messages, start_time=start_time, end_time=end_time
# 		)
# 		self.assertEqual(len(filtered), 1)
# 		self.assertTrue(
# 			start_time <= filtered.iloc[0]["timestamp"] <= end_time
# 			)


# class TestRemoveMessage(unittest.TestCase):

# 	def test_remove_message(self):
# 		"""Test removing a message by node_id."""
# 		messages = pd.DataFrame(
# 			{
# 				"node_id": ["1", "2"], "role": ["user", "assistant"],
# 				"content": ["message1", "message2"],
# 			}
# 		)
# 		updated_messages = MessageUtil.remove_message(messages, "1")
# 		self.assertTrue(updated_messages)


# class TestGetMessageRows(unittest.TestCase):

# 	def test_get_message_rows(self):
# 		"""Test retrieving the last 'n' message rows based on criteria."""
# 		messages = pd.DataFrame(
# 			{
# 				"node_id": ["1", "2", "3"],
# 				"role": ["user", "assistant", "user"],
# 				"sender": ["A", "B", "A"],
# 				"content": ["message1", "message2", "message3"],
# 			}
# 		)
# 		rows = MessageUtil.get_message_rows(
# 			messages, sender="A", role="user", n=2
# 			)
# 		self.assertEqual(len(rows), 2)


# # class TestExtend(unittest.TestCase):

# # def test_extend(self):
# #     """Test extending one DataFrame with another, ensuring no duplicate 'node_id'."""
# #     df1 = pd.DataFrame({
# #         "node_id": ["1"],
# #         "role": ["user"],
# #         "sender": ["test"],
# #         "timestamp": ["2020-01-01T00:00:00"],
# #         "content": ['{"message": "test"}']
# #     })
# #     df2 = pd.DataFrame({
# #         "node_id": ["2"],
# #         "role": ["user"],
# #         "sender": ["test"],
# #         "timestamp": ["2020-01-02T00:00:00"],
# #         "content": ['{"message": "test2"}']
# #     })
# #     combined = MessageUtil.extend(df1, df2)
# #     self.assertEqual(len(combined), 2)


# class TestToMarkdownString(unittest.TestCase):

# 	def test_to_markdown_string(self):
# 		"""Test converting messages to a Markdown-formatted string."""
# 		messages = pd.DataFrame(
# 			{
# 				"node_id": ["1"], "role": ["user"],
# 				"content": [json.dumps({"instruction": "Hello, World!"})],
# 			}
# 		)
# 		markdown_str = MessageUtil.to_markdown_string(messages)
# 		self.assertIn("Hello, World!", markdown_str)


# # class TestSearchKeywords(unittest.TestCase):

# #     def test_search_keywords(self):
# #         """Test filtering DataFrame for rows containing specified keywords."""
# #         messages = pd.DataFrame(
# #             {"node_id": ["1", "2"], "content": ["Hello world", "Goodbye world"]}
# #         )
# #         filtered = MessageUtil.search_keywords(messages, "Hello")
# #         print(filtered)
# #         self.assertEqual(len(filtered), 1)


# # class TestReplaceKeyword(unittest.TestCase):

# #     def test_replace_keyword(self):
# #         """Test replacing a keyword in DataFrame's specified column."""
# #         messages = pd.DataFrame({"content": ["Hello world", "Goodbye world"]})
# #         MessageUtil.replace_keyword(messages, "world", "universe")
# #         self.assertTrue(all(messages["content"].str.contains("universe")))


# # class TestReadCsv(unittest.TestCase):

# # @patch("pandas.read_csv")
# # def test_read_csv(self, mock_read_csv):
# #     """Test reading a CSV file into a DataFrame."""
# #     mock_df = pd.DataFrame(
# #         {"node_id": ["1", "2"], "content": ["Hello, World!", "Goodbye, World!"]}
# #     )
# #     mock_read_csv.return_value = mock_df

# #     df = MessageUtil.read_csv("path/to/nonexistent/file.csv")

# #     mock_read_csv.assert_called_once_with("path/to/nonexistent/file.csv")

# #     self.assertTrue(isinstance(df, pd.DataFrame))
# #     self.assertEqual(len(df), 2)
# #     self.assertEqual(list(df.columns), ["node_id", "content"])


# # class TestReadJson(unittest.TestCase):

# #     @patch("pandas.read_json")
# #     def test_read_json(self, mock_read_json):
# #         """Test reading a JSON file into a DataFrame."""
# #         mock_df = pd.DataFrame(
# #             {"node_id": ["1", "2"], "content": ["JSON Message 1", "JSON Message 2"]}
# #         )
# #         mock_read_json.return_value = mock_df

# #         df = MessageUtil.read_json("path/to/nonexistent/file.json")

# #         mock_read_json.assert_called_once_with("path/to/nonexistent/file.json")

# #         self.assertTrue(isinstance(df, pd.DataFrame))
# #         self.assertEqual(len(df), 2)
# #         self.assertEqual(list(df.columns), ["node_id", "content"])


# # class TestRemoveLastNRows(unittest.TestCase):

# #     def test_remove_last_n_rows(self):
# #         """Test removing the last 'n' rows from a DataFrame."""
# #         messages = pd.DataFrame({"content": ["message1", "message2", "message3"]})
# #         updated = MessageUtil.remove_last_n_rows(messages, 2)
# #         self.assertEqual(len(updated), 1)


# # class TestUpdateRow(unittest.TestCase):

# #     def test_update_row(self):
# #         """Test updating a row's value for a specified column."""
# #         messages = pd.DataFrame(
# #             {"node_id": ["1", "2"], "content": ["message1", "message2"]}
# #         )
# #         success = MessageUtil.update_row(messages, 0, "node_id", "3")
# #         self.assertTrue(success)
# #         self.assertTrue("3" in messages["node_id"].values)


# if __name__ == "__main__":
# 	unittest.main()
