diff --git a/tests/regtest/test_base_chat_interface.py b/tests/regtest/test_base_chat_interface.py
index f7d3db5..8945f9a 100644
--- a/tests/regtest/test_base_chat_interface.py
+++ b/tests/regtest/test_base_chat_interface.py
@@ -2,12 +2,12 @@
 
 import pytest
 
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms import FakeListLLM
 from langchain.schema import BaseMessage, HumanMessage, AIMessage, SystemMessage
 
 from esbmc_ai.ai_models import AIModel
 from esbmc_ai.chat_response import ChatResponse
-from esbmc_ai.base_chat_interface import BaseChatInterface
+from esbmc_ai.chats.base_chat_interface import BaseChatInterface
 from esbmc_ai.config import AIAgentConversation, ChatPromptSettings
 
 
diff --git a/tests/test_base_chat_interface.py b/tests/test_base_chat_interface.py
index 8d79412..42b7b1a 100644
--- a/tests/test_base_chat_interface.py
+++ b/tests/test_base_chat_interface.py
@@ -2,10 +2,10 @@
 
 import pytest
 
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms import FakeListLLM
 from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
 from esbmc_ai.ai_models import AIModel
-from esbmc_ai.base_chat_interface import BaseChatInterface
+from esbmc_ai.chats.base_chat_interface import BaseChatInterface
 from esbmc_ai.chat_response import ChatResponse
 from esbmc_ai.config import AIAgentConversation, ChatPromptSettings
 
diff --git a/tests/test_latest_state_solution_generator.py b/tests/test_latest_state_solution_generator.py
index 61edf4e..679e2d6 100644
--- a/tests/test_latest_state_solution_generator.py
+++ b/tests/test_latest_state_solution_generator.py
@@ -9,7 +9,7 @@
 from esbmc_ai.ai_models import AIModel
 from esbmc_ai.chat_response import ChatResponse
 from esbmc_ai.config import AIAgentConversation, ChatPromptSettings
-from esbmc_ai.latest_state_solution_generator import LatestStateSolutionGenerator
+from esbmc_ai.chats.latest_state_solution_generator import LatestStateSolutionGenerator
 
 
 @pytest.fixture(scope="function")
diff --git a/tests/test_solution.py b/tests/test_solution.py
index c3b6fa0..f5d189f 100644
--- a/tests/test_solution.py
+++ b/tests/test_solution.py
@@ -3,6 +3,10 @@
 import pytest
 from esbmc_ai.solution import Solution, SourceFile
 
+#####################################
+# Solution
+#####################################
+
 
 @pytest.fixture(scope="function")
 def solution() -> Solution:
@@ -44,18 +48,3 @@ def test_apply_line_patch() -> None:
     text = "\n".join(["a", "b", "c", "d", "e", "f", "g"])
     answer = "\n".join(["a", "b", "c", "1", "e", "f", "g"])
     assert SourceFile.apply_line_patch(text, "1", 3, 3) == answer
-
-
-def test_get_patch() -> None:
-    raise NotImplementedError()
-
-
-def test_update_content() -> None:
-    raise NotImplementedError()
-
-
-def test_assign_verifier_output() -> None:
-    raise NotImplementedError()
-
-
-# TODO Add more tests for solution and sourcefile
diff --git a/tests/test_solution_generator.py b/tests/test_solution_generator.py
index be6a209..ff86e7b 100644
--- a/tests/test_solution_generator.py
+++ b/tests/test_solution_generator.py
@@ -6,7 +6,7 @@
 
 from esbmc_ai.ai_models import AIModel
 from esbmc_ai.config import AIAgentConversation, ChatPromptSettings
-from esbmc_ai.solution_generator import SolutionGenerator
+from esbmc_ai.chats.solution_generator import SolutionGenerator
 
 
 @pytest.fixture(scope="function")
diff --git a/tests/test_user_chat.py b/tests/test_user_chat.py
index 3c2eca1..f0198ca 100644
--- a/tests/test_user_chat.py
+++ b/tests/test_user_chat.py
@@ -2,13 +2,13 @@
 
 import pytest
 
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms import FakeListLLM
 from langchain.schema import AIMessage, SystemMessage
 
 from esbmc_ai.ai_models import AIModel
 from esbmc_ai.chat_response import ChatResponse, FinishReason
 from esbmc_ai.config import AIAgentConversation, ChatPromptSettings
-from esbmc_ai.user_chat import UserChat
+from esbmc_ai.chats.user_chat import UserChat
 
 
 @pytest.fixture