finished adding simple mock tests for web_api_testing

Author: Diana Strauss

tests/test_web_api_documentation.py

@@ -3,12 +3,13 @@
 
 from hackingBuddyGPT.usecases import SimpleWebAPITesting
 from hackingBuddyGPT.usecases.web import MinimalWebTesting
-from hackingBuddyGPT.usecases.web_api_testing.simple_openapi_documentation import SimpleWebAPIDocumentationUseCase
+from hackingBuddyGPT.usecases.web_api_testing.simple_openapi_documentation import SimpleWebAPIDocumentationUseCase, \
+    SimpleWebAPIDocumentation
 from hackingBuddyGPT.utils import DbStorage, Console
 from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 
 
-class TestSimpleWebAPITesting(unittest.TestCase):
+class TestSimpleWebAPIDocumentationTest(unittest.TestCase):
 
     @patch('hackingBuddyGPT.utils.openai.openai_lib.OpenAILib')
     def setUp(self, MockOpenAILib):
@@ -18,20 +19,22 @@ def setUp(self, MockOpenAILib):
         console = Console()
 
         log_db.init()
-        self.agent = SimpleWebAPITesting(llm=self.mock_llm)
+        self.agent = SimpleWebAPIDocumentation(llm=self.mock_llm)
         self.agent.init()
         self.simple_api_testing = SimpleWebAPIDocumentationUseCase(
             agent=self.agent,
             log_db=log_db,
             console=console,
-            tag='integration_test_linuxprivesc',
+            tag='webApiDocumentation',
             max_turns=len(self.mock_llm.responses)
         )
         self.simple_api_testing.init()
 
     def test_initial_prompt(self):
         # Test if the initial prompt is set correctly
-        self.assertIn("You're a seasoned API tester tasked with evaluating the REST APIs of a website hosted at https://jsonplaceholder.typicode.com. Your primary objective is to thoroughly explore the APIs, understanding their endpoints, parameters, and responses, and then devise and execute potential exploits. Be diligent in documenting your findings as you navigate through the APIs. Avoid resorting to brute-force tactics. All necessary information can be obtained from the API endpoints themselves. Remember, if you encounter an HTTP method (A string that represents an HTTP method (e.g., 'GET', 'POST', etc.).), promptly submit it as it is of utmost importance.", self.agent._prompt_history[0]['content'])
+        expected_prompt = "You're tasked with documenting the REST APIs of a website hosted at https://jsonplaceholder.typicode.com. Start with an empty OpenAPI specification.\nMaintain meticulousness in documenting your observations as you traverse the APIs."
+
+        self.assertIn(expected_prompt, self.agent._prompt_history[0]['content'])
 
     def test_all_flags_found(self):
         # Mock console.print to suppress output during testing
@@ -62,7 +65,7 @@ def test_perform_round(self, mock_perf_counter):
         result = self.agent.perform_round(1)
 
         # Assertions
-        self.assertFalse(result)  # No flags found in this round
+        self.assertTrue(result)
 
         # Check if the LLM was called with the correct parameters
         mock_create_with_completion = self.agent.llm.instructor.chat.completions.create_with_completion

tests/test_web_api_testing.py

@@ -1,42 +1,78 @@
-import argparse
 import unittest
-from hackingBuddyGPT.usecases.base import use_cases
+from unittest.mock import MagicMock, patch
+
+from hackingBuddyGPT.usecases import SimpleWebAPITesting
+from hackingBuddyGPT.usecases.web import MinimalWebTesting
+from hackingBuddyGPT.usecases.web_api_testing.simple_openapi_documentation import SimpleWebAPIDocumentationUseCase
 from hackingBuddyGPT.usecases.web_api_testing.simple_web_api_testing import SimpleWebAPITestingUseCase
 from hackingBuddyGPT.utils import DbStorage, Console
+from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 
 
+class TestSimpleWebAPITestingTest(unittest.TestCase):
 
-class WebAPITestingTestCase(unittest.TestCase):
-    def test_simple_web_api_testing(self):
-
+    @patch('hackingBuddyGPT.utils.openai.openai_lib.OpenAILib')
+    def setUp(self, MockOpenAILib):
+        # Mock the OpenAILib instance
+        self.mock_llm = MockOpenAILib.return_value
         log_db = DbStorage(':memory:')
         console = Console()
 
         log_db.init()
-        parser = argparse.ArgumentParser()
-        subparser = parser.add_subparsers(required=True)
-        for name, use_case in use_cases.items():
-            use_case.build_parser(subparser.add_parser(
-                name=use_case.name,
-                help=use_case.description
-            ))
-
-        parsed = parser.parse_args(["SimpleWebAPITesting"])
-        instance = parsed.use_case(parsed)
-
-        agent = instance.agent
-        simple_web_api_testing = SimpleWebAPITestingUseCase(
-            agent=agent,
+        self.agent = SimpleWebAPITesting(llm=self.mock_llm)
+        self.agent.init()
+        self.simple_api_testing = SimpleWebAPITestingUseCase(
+            agent=self.agent,
             log_db=log_db,
             console=console,
-            tag='web_api_testing',
-            max_turns=20
+            tag='integration_test_linuxprivesc',
+            max_turns=len(self.mock_llm.responses)
         )
+        self.simple_api_testing.init()
+
+    def test_initial_prompt(self):
+        # Test if the initial prompt is set correctly
+        self.assertIn("You're a seasoned API tester tasked with evaluating the REST APIs of a website hosted at https://jsonplaceholder.typicode.com. Your primary objective is to thoroughly explore the APIs, understanding their endpoints, parameters, and responses, and then devise and execute potential exploits. Be diligent in documenting your findings as you navigate through the APIs. Avoid resorting to brute-force tactics. All necessary information can be obtained from the API endpoints themselves. Remember, if you encounter an HTTP method (A string that represents an HTTP method (e.g., 'GET', 'POST', etc.).), promptly submit it as it is of utmost importance.", self.agent._prompt_history[0]['content'])
+
+    def test_all_flags_found(self):
+        # Mock console.print to suppress output during testing
+        with patch('rich.console.Console.print'):
+            self.agent.all_http_methods_found()
+            self.assertFalse(self.agent.all_http_methods_found())
+
+    @patch('time.perf_counter', side_effect=[1, 2])  # Mocking perf_counter for consistent timing
+    def test_perform_round(self, mock_perf_counter):
+        # Prepare mock responses
+        mock_response = MagicMock()
+        mock_completion = MagicMock()
+
+        # Setup completion response with mocked data
+        mock_completion.choices[0].message.content = "Mocked LLM response"
+        mock_completion.choices[0].message.tool_calls = [MagicMock(id="tool_call_1")]
+        mock_completion.usage.prompt_tokens = 10
+        mock_completion.usage.completion_tokens = 20
+
+        # Mock the OpenAI LLM response
+        self.agent.llm.instructor.chat.completions.create_with_completion.return_value = (
+        mock_response, mock_completion)
+
+        # Mock the tool execution result
+        mock_response.execute.return_value = "Mocked tool execution result"
+
+        # Perform the round
+        result = self.agent.perform_round(1)
+
+        # Assertions
+        self.assertFalse(result)  # No flags found in this round
+
+        # Check if the LLM was called with the correct parameters
+        mock_create_with_completion = self.agent.llm.instructor.chat.completions.create_with_completion
 
-        simple_web_api_testing.init()
-        result = simple_web_api_testing.run()
-        # TODO: find condition for testing
+        # if it can be called multiple times, use assert_called
+        self.assertEqual( 2, mock_create_with_completion.call_count)
 
+        # Check if the prompt history was updated correctly
+        self.assertEqual(5, len(self.agent._prompt_history))  # Initial message + LLM response + tool message
 
 if __name__ == '__main__':
     unittest.main()