From f1831bfccecb9c546f7b35773ab0fb9c97559848 Mon Sep 17 00:00:00 2001
From: Stephen Lincoln <stephen.lincoln@attackiq.com>
Date: Mon, 4 Nov 2024 21:12:42 -0500
Subject: [PATCH] Fixed asyncio tests

---
 tests/test_llm_components.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/tests/test_llm_components.py b/tests/test_llm_components.py
index e3540f7..1bb465d 100644
--- a/tests/test_llm_components.py
+++ b/tests/test_llm_components.py
@@ -18,7 +18,17 @@ def invoke(self, *args, **kwargs):
         return "Mocked LLM response"
 
     async def ainvoke(self, *args, **kwargs):
-        return "Mocked async LLM response"
+        return """title: Detect PowerShell Execution
+description: Detects PowerShell execution
+status: test
+author: MockLLM
+logsource:
+    category: process_creation
+    product: windows
+detection:
+    selection:
+        CommandLine|contains: 'powershell'
+    condition: selection"""
 
     def generate_prompt(self, *args, **kwargs):
         return "Mocked generate_prompt response"
@@ -94,7 +104,7 @@ def test_sigma_toolkit():
 
 @pytest.mark.asyncio
 async def test_create_sigma_rule_tool(mock_openai_create, mock_vector_store):
-    tool = CreateSigmaRuleVectorStoreTool(sigmadb=mock_vector_store, llm=Mock())
+    tool = CreateSigmaRuleVectorStoreTool(sigmadb=mock_vector_store, llm=MockLLM())
     result = await tool._arun("Create a Sigma rule for detecting PowerShell execution")
     assert isinstance(result, str)
     assert "title:" in result.lower()
@@ -110,13 +120,13 @@ async def test_translate_sigma_rule_tool(mock_openai_create):
 
 @pytest.mark.asyncio
 async def test_find_sigma_rule_tool(mock_openai_create, mock_vector_store):
-    tool = FindSigmaRuleTool(sigmadb=mock_vector_store, llm=Mock())
+    tool = FindSigmaRuleTool(sigmadb=mock_vector_store, llm=MockLLM())
     result = await tool._arun("Find a rule for detecting mimikatz")
     assert isinstance(result, str)
 
 @pytest.mark.asyncio
 async def test_query_to_sigma_rule_tool(mock_openai_create):
-    tool = QueryToSigmaRuleTool(llm=Mock())
+    tool = QueryToSigmaRuleTool(llm=MockLLM())
     result = await tool._arun(
         query="process_name=powershell.exe",
         backend="splunk"