From a3375206bfc1662cc51cb2a96fe9edaf4bef52ce Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 11:43:56 +0200 Subject: [PATCH 1/9] fixed web_api_documentation test and removed unnecessary imports --- .../web_api_testing/prompt_engineer.py | 18 +++++++++--------- tests/test_web_api_documentation.py | 6 +----- tests/test_web_api_testing.py | 4 ---- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py b/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py index b3efbf5..cd1bb40 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py @@ -1,6 +1,5 @@ import nltk from nltk.tokenize import word_tokenize -from nltk.corpus import stopwords from instructor.retry import InstructorRetryException @@ -37,13 +36,11 @@ def __init__(self, strategy, llm_handler, history, schemas, response_handler): self.found_endpoints = ["/"] self.endpoint_methods = {} self.endpoint_found_methods = {} - model_name = "en_core_web_sm" - # Check if the models are already installed nltk.download('punkt') nltk.download('stopwords') self._prompt_history = history - self.prompt = self._prompt_history + self.prompt = {self.round: {"content": "initial_prompt"}} self.previous_prompt = self._prompt_history[self.round]["content"] self.schemas = schemas @@ -77,10 +74,6 @@ def generate_prompt(self, doc=False): self.round = self.round +1 return self._prompt_history - - - - def in_context_learning(self, doc=False, hint=""): """ Generates a prompt for in-context learning. @@ -91,7 +84,14 @@ def in_context_learning(self, doc=False, hint=""): Returns: str: The generated prompt. """ - return str("\n".join(self._prompt_history[self.round]["content"] + [self.prompt])) + history_content = [entry["content"] for entry in self._prompt_history] + prompt_content = self.prompt.get(self.round, {}).get("content", "") + + # Add hint if provided + if hint: + prompt_content += f"\n{hint}" + + return "\n".join(history_content + [prompt_content]) def get_http_action_template(self, method): """Helper to construct a consistent HTTP action description.""" diff --git a/tests/test_web_api_documentation.py b/tests/test_web_api_documentation.py index 3e54ac0..a47b2b1 100644 --- a/tests/test_web_api_documentation.py +++ b/tests/test_web_api_documentation.py @@ -1,12 +1,8 @@ import unittest 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, \ SimpleWebAPIDocumentation from hackingBuddyGPT.utils import DbStorage, Console -from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib class TestSimpleWebAPIDocumentationTest(unittest.TestCase): @@ -40,7 +36,7 @@ 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()) + self.assertTrue(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): diff --git a/tests/test_web_api_testing.py b/tests/test_web_api_testing.py index cc151aa..e264cfc 100644 --- a/tests/test_web_api_testing.py +++ b/tests/test_web_api_testing.py @@ -1,12 +1,8 @@ import unittest 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): From 947c8a7914c4b8e25bda97d35b11215c7cace8e9 Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 11:44:17 +0200 Subject: [PATCH 2/9] Added test for prompt engineer --- tests/test_prompt_engineer.py | 64 +++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 tests/test_prompt_engineer.py diff --git a/tests/test_prompt_engineer.py b/tests/test_prompt_engineer.py new file mode 100644 index 0000000..f8b9e44 --- /dev/null +++ b/tests/test_prompt_engineer.py @@ -0,0 +1,64 @@ +import unittest +from unittest.mock import MagicMock +from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptStrategy, PromptEngineer + + +class TestPromptEngineer(unittest.TestCase): + def setUp(self): + self.strategy = PromptStrategy.IN_CONTEXT + self.llm_handler = MagicMock() + self.history = [{"content": "initial_prompt", "role": "system"}] + self.schemas = MagicMock() + self.response_handler = MagicMock() + self.prompt_engineer = PromptEngineer( + self.strategy, self.llm_handler, self.history, self.schemas, self.response_handler + ) + def test_token_count(self): + text = "This is a sample text with several words." + count = self.prompt_engineer.token_count(text) + self.assertEqual(8, count) + def test_check_prompt(self): + self.response_handler.get_response_for_prompt = MagicMock(return_value="shortened_prompt") + prompt = self.prompt_engineer.check_prompt("previous_prompt", + ["step1", "step2", "step3", "step4", "step5", "step6"], max_tokens=5) + self.assertEqual(prompt, "shortened_prompt") + + def test_in_context_learning_no_hint(self): + expected_prompt = "initial_prompt\ninitial_prompt" + actual_prompt = self.prompt_engineer.in_context_learning() + self.assertEqual(expected_prompt, actual_prompt) + + def test_in_context_learning_with_hint(self): + hint = "This is a hint." + expected_prompt = "initial_prompt\ninitial_prompt\nThis is a hint." + actual_prompt = self.prompt_engineer.in_context_learning(hint=hint) + self.assertEqual(expected_prompt, actual_prompt) + + def test_in_context_learning_with_doc_and_hint(self): + hint = "This is another hint." + expected_prompt = "initial_prompt\ninitial_prompt\nThis is another hint." + actual_prompt = self.prompt_engineer.in_context_learning(doc=True, hint=hint) + self.assertEqual(expected_prompt, actual_prompt) + def test_generate_prompt_chain_of_thought(self): + self.prompt_engineer.strategy = PromptStrategy.CHAIN_OF_THOUGHT + self.response_handler.get_response_for_prompt = MagicMock(return_value="response_text") + self.prompt_engineer.evaluate_response = MagicMock(return_value=True) + + prompt_history = self.prompt_engineer.generate_prompt() + + self.assertEqual( 2, len(prompt_history)) + + def test_generate_prompt_tree_of_thought(self): + self.prompt_engineer.strategy = PromptStrategy.TREE_OF_THOUGHT + self.response_handler.get_response_for_prompt = MagicMock(return_value="response_text") + self.prompt_engineer.evaluate_response = MagicMock(return_value=True) + + prompt_history = self.prompt_engineer.generate_prompt() + + self.assertEqual(len(prompt_history), 2) + + + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file From 16405382971e2db20d375e9803b9dae079de1c32 Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 12:00:12 +0200 Subject: [PATCH 3/9] Added optional dependencies to .toml file for testing, instructions were also added to README.md --- README.md | 3 +++ pyproject.toml | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/README.md b/README.md index 99f020d..2c536c6 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,9 @@ wintermute.py: error: the following arguments are required: {linux_privesc,windo # start wintermute, i.e., attack the configured virtual machine $ python wintermute.py minimal_linux_privesc + +# install dependencies for testing if you want to run the tests +$ pip install .[testing] ~~~ ## Publications about hackingBuddyGPT diff --git a/pyproject.toml b/pyproject.toml index 6e9084b..8133e90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,11 @@ pythonpath = "src" addopts = [ "--import-mode=importlib", ] +[project.optional-dependencies] +testing = [ + 'pytest', + 'pytest-mock' +] [project.scripts] wintermute = "hackingBuddyGPT.cli.wintermute:main" From 86cf6484a1c8af9a5c59e95547759b4939e36b43 Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 13:51:00 +0200 Subject: [PATCH 4/9] Changed name of documentation_handler of web_api as there were other classes with that name already, it is now classed OpenAPISpecificationManager; Tests for this were also added --- .../simple_openapi_documentation.py | 4 +- .../web_api_testing/utils/__init__.py | 2 +- ...er.py => openAPI_specification_manager.py} | 8 +-- tests/test_openAPI_specification_manager.py | 53 +++++++++++++++++++ 4 files changed, 60 insertions(+), 7 deletions(-) rename src/hackingBuddyGPT/usecases/web_api_testing/utils/{documentation_handler.py => openAPI_specification_manager.py} (96%) create mode 100644 tests/test_openAPI_specification_manager.py diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py index 84589cf..be92989 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py @@ -9,7 +9,7 @@ from hackingBuddyGPT.capabilities.http_request import HTTPRequest from hackingBuddyGPT.capabilities.record_note import RecordNote from hackingBuddyGPT.usecases.agents import Agent -from hackingBuddyGPT.usecases.web_api_testing.utils.documentation_handler import DocumentationHandler +from hackingBuddyGPT.usecases.web_api_testing.utils.openAPI_specification_manager import OpenAPISpecificationManager from hackingBuddyGPT.usecases.web_api_testing.utils.llm_handler import LLMHandler from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy from hackingBuddyGPT.usecases.web_api_testing.utils.response_handler import ResponseHandler @@ -52,7 +52,7 @@ def init(self): self.llm_handler = LLMHandler(self.llm, self._capabilities) self.response_handler = ResponseHandler(self.llm_handler) self._setup_initial_prompt() - self.documentation_handler = DocumentationHandler(self.llm_handler, self.response_handler) + self.documentation_handler = OpenAPISpecificationManager(self.llm_handler, self.response_handler) def _setup_capabilities(self): notes = self._context["notes"] diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py b/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py index 0b7f74a..388829a 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py @@ -1,4 +1,4 @@ -from .documentation_handler import DocumentationHandler +from .openAPI_specification_manager import OpenAPISpecificationManager from .llm_handler import LLMHandler from .response_handler import ResponseHandler from .openapi_parser import OpenAPISpecificationParser diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/utils/documentation_handler.py b/src/hackingBuddyGPT/usecases/web_api_testing/utils/openAPI_specification_manager.py similarity index 96% rename from src/hackingBuddyGPT/usecases/web_api_testing/utils/documentation_handler.py rename to src/hackingBuddyGPT/usecases/web_api_testing/utils/openAPI_specification_manager.py index d03d023..bdfc2e7 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/utils/documentation_handler.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/utils/openAPI_specification_manager.py @@ -3,7 +3,7 @@ from datetime import datetime from hackingBuddyGPT.capabilities.yamlFile import YAMLFile -class DocumentationHandler: +class OpenAPISpecificationManager: """ Handles the generation and updating of an OpenAPI specification document based on dynamic API responses. @@ -51,7 +51,7 @@ def __init__(self, llm_handler, response_handler): "yaml": YAMLFile() } - def partial_match(self, element, string_list): + def is_partial_match(self, element, string_list): return any(element in string or string in element for string in string_list) def update_openapi_spec(self, resp, result): @@ -66,7 +66,7 @@ def update_openapi_spec(self, resp, result): if request.__class__.__name__ == 'RecordNote': # TODO: check why isinstance does not work self.check_openapi_spec(resp) - if request.__class__.__name__ == 'HTTPRequest': + elif request.__class__.__name__ == 'HTTPRequest': path = request.path method = request.method print(f'method: {method}') @@ -107,7 +107,7 @@ def update_openapi_spec(self, resp, result): if '1' not in path and x != "": endpoint_methods[path].append(method) - elif self.partial_match(x, endpoints.keys()): + elif self.is_partial_match(x, endpoints.keys()): path = f"/{x}" print(f'endpoint methods = {endpoint_methods}') print(f'new path:{path}') diff --git a/tests/test_openAPI_specification_manager.py b/tests/test_openAPI_specification_manager.py new file mode 100644 index 0000000..35b5dc2 --- /dev/null +++ b/tests/test_openAPI_specification_manager.py @@ -0,0 +1,53 @@ +import unittest +from unittest.mock import MagicMock, patch + +from hackingBuddyGPT.capabilities.http_request import HTTPRequest +from hackingBuddyGPT.usecases.web_api_testing.utils import OpenAPISpecificationManager + + +class TestSpecificationHandler(unittest.TestCase): + def setUp(self): + self.llm_handler = MagicMock() + self.response_handler = MagicMock() + self.doc_handler = OpenAPISpecificationManager(self.llm_handler, self.response_handler) + + @patch('os.makedirs') + @patch('builtins.open') + def test_write_openapi_to_yaml(self, mock_open, mock_makedirs): + self.doc_handler.write_openapi_to_yaml() + mock_makedirs.assert_called_once_with(self.doc_handler.file_path, exist_ok=True) + mock_open.assert_called_once_with(self.doc_handler.file, 'w') + + # Create a mock HTTPRequest object + response_mock = MagicMock() + response_mock.action = HTTPRequest( + host="https://jsonplaceholder.typicode.com", + follow_redirects=False, + use_cookie_jar=True + ) + response_mock.action.method = "GET" + response_mock.action.path = "/test" + + result = '{"key": "value"}' + + self.response_handler.parse_http_response_to_openapi_example = MagicMock( + return_value=({}, "#/components/schemas/TestSchema", self.doc_handler.openapi_spec) + ) + + endpoints = self.doc_handler.update_openapi_spec(response_mock, result) + + self.assertIn("/test", self.doc_handler.openapi_spec["endpoints"]) + self.assertIn("get", self.doc_handler.openapi_spec["endpoints"]["/test"]) + self.assertEqual(self.doc_handler.openapi_spec["endpoints"]["/test"]["get"]["summary"], + "GET operation on /test") + self.assertEqual(endpoints, ["/test"]) + + + def test_partial_match(self): + string_list = ["test_endpoint", "another_endpoint"] + self.assertTrue(self.doc_handler.is_partial_match("test", string_list)) + self.assertFalse(self.doc_handler.is_partial_match("not_in_list", string_list)) + + +if __name__ == "__main__": + unittest.main() From 44af8180aa7780e89466a7e4fe92073acb6cd2e4 Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 14:27:57 +0200 Subject: [PATCH 5/9] Added tests for llm_handler and response_handler --- .../simple_openapi_documentation.py | 2 +- .../web_api_testing/utils/__init__.py | 2 +- ...er.py => openapi_specification_manager.py} | 0 .../web_api_testing/utils/response_handler.py | 11 +- tests/test_llm_handler.py | 61 +++++++++ tests/test_response_handler.py | 117 ++++++++++++++++++ tests/test_web_api_documentation.py | 2 +- tests/test_web_api_testing.py | 2 +- 8 files changed, 189 insertions(+), 8 deletions(-) rename src/hackingBuddyGPT/usecases/web_api_testing/utils/{openAPI_specification_manager.py => openapi_specification_manager.py} (100%) create mode 100644 tests/test_llm_handler.py create mode 100644 tests/test_response_handler.py diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py index be92989..c1e9c41 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py @@ -9,7 +9,7 @@ from hackingBuddyGPT.capabilities.http_request import HTTPRequest from hackingBuddyGPT.capabilities.record_note import RecordNote from hackingBuddyGPT.usecases.agents import Agent -from hackingBuddyGPT.usecases.web_api_testing.utils.openAPI_specification_manager import OpenAPISpecificationManager +from hackingBuddyGPT.usecases.web_api_testing.utils.openapi_specification_manager import OpenAPISpecificationManager from hackingBuddyGPT.usecases.web_api_testing.utils.llm_handler import LLMHandler from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy from hackingBuddyGPT.usecases.web_api_testing.utils.response_handler import ResponseHandler diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py b/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py index 388829a..a856540 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/utils/__init__.py @@ -1,4 +1,4 @@ -from .openAPI_specification_manager import OpenAPISpecificationManager +from .openapi_specification_manager import OpenAPISpecificationManager from .llm_handler import LLMHandler from .response_handler import ResponseHandler from .openapi_parser import OpenAPISpecificationParser diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/utils/openAPI_specification_manager.py b/src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_specification_manager.py similarity index 100% rename from src/hackingBuddyGPT/usecases/web_api_testing/utils/openAPI_specification_manager.py rename to src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_specification_manager.py diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py b/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py index 47e47c7..4b2dcae 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py @@ -50,11 +50,14 @@ def parse_http_status_line(self, status_line): """ if status_line == "Not a valid HTTP method": return status_line - if status_line and " " in status_line: - protocol, status_code, status_message = status_line.split(' ', 2) - status_message = status_message.split("\r\n")[0] + + # Regular expression to match valid HTTP status lines + match = re.match(r'^(HTTP/\d\.\d) (\d{3}) (.*)$', status_line) + if match: + protocol, status_code, status_message = match.groups() return f'{status_code} {status_message}' - raise ValueError("Invalid HTTP status line") + else: + raise ValueError("Invalid HTTP status line") def extract_response_example(self, html_content): """ diff --git a/tests/test_llm_handler.py b/tests/test_llm_handler.py new file mode 100644 index 0000000..9b209d2 --- /dev/null +++ b/tests/test_llm_handler.py @@ -0,0 +1,61 @@ +import unittest +from unittest.mock import MagicMock, patch +from hackingBuddyGPT.capabilities.capability import capabilities_to_action_model +from hackingBuddyGPT.usecases.web_api_testing.utils import LLMHandler + + +class TestLLMHandler(unittest.TestCase): + def setUp(self): + self.llm_mock = MagicMock() + self.capabilities = {'cap1': MagicMock(), 'cap2': MagicMock()} + self.llm_handler = LLMHandler(self.llm_mock, self.capabilities) + + '''@patch('hackingBuddyGPT.usecases.web_api_testing.utils.capabilities_to_action_model') + def test_call_llm(self, mock_capabilities_to_action_model): + prompt = [{'role': 'user', 'content': 'Hello, LLM!'}] + response_mock = MagicMock() + self.llm_mock.instructor.chat.completions.create_with_completion.return_value = response_mock + + # Mock the capabilities_to_action_model to return a dummy Pydantic model + mock_model = MagicMock() + mock_capabilities_to_action_model.return_value = mock_model + + response = self.llm_handler.call_llm(prompt) + + self.llm_mock.instructor.chat.completions.create_with_completion.assert_called_once_with( + model=self.llm_mock.model, + messages=prompt, + response_model=mock_model + ) + self.assertEqual(response, response_mock)''' + def test_add_created_object(self): + created_object = MagicMock() + object_type = 'test_type' + + self.llm_handler.add_created_object(created_object, object_type) + + self.assertIn(object_type, self.llm_handler.created_objects) + self.assertIn(created_object, self.llm_handler.created_objects[object_type]) + + def test_add_created_object_limit(self): + created_object = MagicMock() + object_type = 'test_type' + + for _ in range(8): # Exceed the limit of 7 objects + self.llm_handler.add_created_object(created_object, object_type) + + self.assertEqual(len(self.llm_handler.created_objects[object_type]), 7) + + def test_get_created_objects(self): + created_object = MagicMock() + object_type = 'test_type' + self.llm_handler.add_created_object(created_object, object_type) + + created_objects = self.llm_handler.get_created_objects() + + self.assertIn(object_type, created_objects) + self.assertIn(created_object, created_objects[object_type]) + self.assertEqual(created_objects, self.llm_handler.created_objects) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_response_handler.py b/tests/test_response_handler.py new file mode 100644 index 0000000..d80ce55 --- /dev/null +++ b/tests/test_response_handler.py @@ -0,0 +1,117 @@ +import unittest +from unittest.mock import MagicMock, patch +from bs4 import BeautifulSoup +import json +from hackingBuddyGPT.usecases.web_api_testing.utils import ResponseHandler + +class TestResponseHandler(unittest.TestCase): + def setUp(self): + self.llm_handler_mock = MagicMock() + self.response_handler = ResponseHandler(self.llm_handler_mock) + + def test_get_response_for_prompt(self): + prompt = "Test prompt" + response_mock = MagicMock() + response_mock.execute.return_value = "Response text" + self.llm_handler_mock.call_llm.return_value = (response_mock, MagicMock()) + + response_text = self.response_handler.get_response_for_prompt(prompt) + + self.llm_handler_mock.call_llm.assert_called_once_with([{"role": "user", "content": [{"type": "text", "text": prompt}]}]) + self.assertEqual(response_text, "Response text") + + def test_parse_http_status_line_valid(self): + status_line = "HTTP/1.1 200 OK" + result = self.response_handler.parse_http_status_line(status_line) + self.assertEqual(result, "200 OK") + + def test_parse_http_status_line_invalid(self): + status_line = "Invalid status line" + with self.assertRaises(ValueError): + self.response_handler.parse_http_status_line(status_line) + + def test_extract_response_example(self): + html_content = """ + + + {"example": "test"} + {"key": "value"} + + + """ + result = self.response_handler.extract_response_example(html_content) + self.assertEqual(result, {"key": "value"}) + + def test_extract_response_example_invalid(self): + html_content = "No code tags" + result = self.response_handler.extract_response_example(html_content) + self.assertIsNone(result) + + @patch('hackingBuddyGPT.usecases.web_api_testing.utils.ResponseHandler.parse_http_response_to_schema') + def test_parse_http_response_to_openapi_example(self, mock_parse_http_response_to_schema): + openapi_spec = { + "components": {"schemas": {}} + } + http_response = "HTTP/1.1 200 OK\r\n\r\n{\"id\": 1, \"name\": \"test\"}" + path = "/test" + method = "GET" + + mock_parse_http_response_to_schema.return_value = ("#/components/schemas/Test", "Test", openapi_spec) + + entry_dict, reference, updated_spec = self.response_handler.parse_http_response_to_openapi_example(openapi_spec, http_response, path, method) + + self.assertEqual(reference, "#/components/schemas/Test") + self.assertEqual(updated_spec, openapi_spec) + self.assertIn("test", entry_dict) + + def test_extract_description(self): + note = MagicMock() + note.action.content = "Test description" + description = self.response_handler.extract_description(note) + self.assertEqual(description, "Test description") + + @patch('hackingBuddyGPT.usecases.web_api_testing.utils.ResponseHandler.extract_keys') + def test_parse_http_response_to_schema(self, mock_extract_keys): + openapi_spec = { + "components": {"schemas": {}} + } + body_dict = {"id": 1, "name": "test"} + path = "/tests" + + mock_extract_keys.side_effect = lambda key, value, properties: {**properties, key: {"type": type(value).__name__, "example": value}} + + reference, object_name, updated_spec = self.response_handler.parse_http_response_to_schema(openapi_spec, body_dict, path) + + self.assertEqual(reference, "#/components/schemas/Test") + self.assertEqual(object_name, "Test") + self.assertIn("Test", updated_spec["components"]["schemas"]) + self.assertIn("id", updated_spec["components"]["schemas"]["Test"]["properties"]) + self.assertIn("name", updated_spec["components"]["schemas"]["Test"]["properties"]) + + @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data='yaml_content') + def test_read_yaml_to_string(self, mock_open): + filepath = "test.yaml" + result = self.response_handler.read_yaml_to_string(filepath) + mock_open.assert_called_once_with(filepath, 'r') + self.assertEqual(result, 'yaml_content') + + def test_read_yaml_to_string_file_not_found(self): + filepath = "nonexistent.yaml" + result = self.response_handler.read_yaml_to_string(filepath) + self.assertIsNone(result) + + def test_extract_endpoints(self): + note = "1. GET /test\n" + result = self.response_handler.extract_endpoints(note) + self.assertEqual( {'/test': ['GET']}, result) + + def test_extract_keys(self): + key = "name" + value = "test" + properties_dict = {} + result = self.response_handler.extract_keys(key, value, properties_dict) + self.assertIn(key, result) + self.assertEqual(result[key], {"type": "str", "example": "test"}) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_web_api_documentation.py b/tests/test_web_api_documentation.py index a47b2b1..8ef076a 100644 --- a/tests/test_web_api_documentation.py +++ b/tests/test_web_api_documentation.py @@ -55,7 +55,7 @@ def test_perform_round(self, mock_perf_counter): mock_response, mock_completion) # Mock the tool execution result - mock_response.execute.return_value = "Mocked tool execution result" + mock_response.execute.return_value = "HTTP/1.1 200 OK" # Perform the round result = self.agent.perform_round(1) diff --git a/tests/test_web_api_testing.py b/tests/test_web_api_testing.py index e264cfc..aa4d5da 100644 --- a/tests/test_web_api_testing.py +++ b/tests/test_web_api_testing.py @@ -53,7 +53,7 @@ def test_perform_round(self, mock_perf_counter): mock_response, mock_completion) # Mock the tool execution result - mock_response.execute.return_value = "Mocked tool execution result" + mock_response.execute.return_value = "HTTP/1.1 200 OK" # Perform the round result = self.agent.perform_round(1) From 9bdd6bdef4ad92524cd1426ce0a9095e6eadb128 Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 14:33:27 +0200 Subject: [PATCH 6/9] Added tests for openapi converter and parser --- tests/test_openapi_converter.py | 87 +++++++++++ tests/test_openapi_parser.py | 256 ++++++++++++++++++++++++++++++++ 2 files changed, 343 insertions(+) create mode 100644 tests/test_openapi_converter.py create mode 100644 tests/test_openapi_parser.py diff --git a/tests/test_openapi_converter.py b/tests/test_openapi_converter.py new file mode 100644 index 0000000..67e1741 --- /dev/null +++ b/tests/test_openapi_converter.py @@ -0,0 +1,87 @@ +import unittest +from unittest.mock import patch, mock_open, MagicMock +import os +import yaml +import json + +from hackingBuddyGPT.usecases.utils.openapi_converter import OpenAPISpecificationConverter + + +class TestOpenAPISpecificationConverter(unittest.TestCase): + def setUp(self): + self.converter = OpenAPISpecificationConverter("base_directory") + + @patch("os.makedirs") + @patch("builtins.open", new_callable=mock_open, read_data="yaml_content") + @patch("yaml.safe_load", return_value={"key": "value"}) + @patch("json.dump") + def test_convert_file_yaml_to_json(self, mock_json_dump, mock_yaml_safe_load, mock_open_file, mock_makedirs): + input_filepath = "input.yaml" + output_directory = "json" + input_type = "yaml" + output_type = "json" + expected_output_path = os.path.join("base_directory", output_directory, "input.json") + + result = self.converter.convert_file(input_filepath, output_directory, input_type, output_type) + + mock_open_file.assert_any_call(input_filepath, 'r') + mock_yaml_safe_load.assert_called_once() + mock_open_file.assert_any_call(expected_output_path, 'w') + mock_json_dump.assert_called_once_with({"key": "value"}, mock_open_file(), indent=2) + mock_makedirs.assert_called_once_with(os.path.join("base_directory", output_directory), exist_ok=True) + self.assertEqual(result, expected_output_path) + + @patch("os.makedirs") + @patch("builtins.open", new_callable=mock_open, read_data='{"key": "value"}') + @patch("json.load", return_value={"key": "value"}) + @patch("yaml.dump") + def test_convert_file_json_to_yaml(self, mock_yaml_dump, mock_json_load, mock_open_file, mock_makedirs): + input_filepath = "input.json" + output_directory = "yaml" + input_type = "json" + output_type = "yaml" + expected_output_path = os.path.join("base_directory", output_directory, "input.yaml") + + result = self.converter.convert_file(input_filepath, output_directory, input_type, output_type) + + mock_open_file.assert_any_call(input_filepath, 'r') + mock_json_load.assert_called_once() + mock_open_file.assert_any_call(expected_output_path, 'w') + mock_yaml_dump.assert_called_once_with({"key": "value"}, mock_open_file(), allow_unicode=True, default_flow_style=False) + mock_makedirs.assert_called_once_with(os.path.join("base_directory", output_directory), exist_ok=True) + self.assertEqual(result, expected_output_path) + + @patch("os.makedirs") + @patch("builtins.open", new_callable=mock_open, read_data="yaml_content") + @patch("yaml.safe_load", side_effect=Exception("YAML error")) + def test_convert_file_yaml_to_json_error(self, mock_yaml_safe_load, mock_open_file, mock_makedirs): + input_filepath = "input.yaml" + output_directory = "json" + input_type = "yaml" + output_type = "json" + + result = self.converter.convert_file(input_filepath, output_directory, input_type, output_type) + + mock_open_file.assert_any_call(input_filepath, 'r') + mock_yaml_safe_load.assert_called_once() + mock_makedirs.assert_called_once_with(os.path.join("base_directory", output_directory), exist_ok=True) + self.assertIsNone(result) + + @patch("os.makedirs") + @patch("builtins.open", new_callable=mock_open, read_data='{"key": "value"}') + @patch("json.load", side_effect=Exception("JSON error")) + def test_convert_file_json_to_yaml_error(self, mock_json_load, mock_open_file, mock_makedirs): + input_filepath = "input.json" + output_directory = "yaml" + input_type = "json" + output_type = "yaml" + + result = self.converter.convert_file(input_filepath, output_directory, input_type, output_type) + + mock_open_file.assert_any_call(input_filepath, 'r') + mock_json_load.assert_called_once() + mock_makedirs.assert_called_once_with(os.path.join("base_directory", output_directory), exist_ok=True) + self.assertIsNone(result) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_openapi_parser.py b/tests/test_openapi_parser.py new file mode 100644 index 0000000..0d52251 --- /dev/null +++ b/tests/test_openapi_parser.py @@ -0,0 +1,256 @@ +import unittest +from unittest.mock import patch, mock_open +import yaml +from hackingBuddyGPT.usecases.web_api_testing.utils import OpenAPISpecificationParser + +class TestOpenAPISpecificationParser(unittest.TestCase): + def setUp(self): + self.filepath = "dummy_path.yaml" + self.yaml_content = """ + openapi: 3.0.0 + info: + title: Sample API + version: 1.0.0 + servers: + - url: https://api.example.com + - url: https://staging.api.example.com + paths: + /pets: + get: + summary: List all pets + responses: + '200': + description: A paged array of pets + post: + summary: Create a pet + responses: + '200': + description: Pet created + /pets/{petId}: + get: + summary: Info for a specific pet + responses: + '200': + description: Expected response to a valid request + """ + + @patch("builtins.open", new_callable=mock_open, read_data="") + @patch("yaml.safe_load", return_value=yaml.safe_load(""" + openapi: 3.0.0 + info: + title: Sample API + version: 1.0.0 + servers: + - url: https://api.example.com + - url: https://staging.api.example.com + paths: + /pets: + get: + summary: List all pets + responses: + '200': + description: A paged array of pets + post: + summary: Create a pet + responses: + '200': + description: Pet created + /pets/{petId}: + get: + summary: Info for a specific pet + responses: + '200': + description: Expected response to a valid request + """)) + def test_load_yaml(self, mock_yaml_load, mock_open_file): + parser = OpenAPISpecificationParser(self.filepath) + self.assertEqual(parser.api_data['info']['title'], "Sample API") + self.assertEqual(parser.api_data['info']['version'], "1.0.0") + self.assertEqual(len(parser.api_data['servers']), 2) + + @patch("builtins.open", new_callable=mock_open, read_data="") + @patch("yaml.safe_load", return_value=yaml.safe_load(""" + openapi: 3.0.0 + info: + title: Sample API + version: 1.0.0 + servers: + - url: https://api.example.com + - url: https://staging.api.example.com + paths: + /pets: + get: + summary: List all pets + responses: + '200': + description: A paged array of pets + post: + summary: Create a pet + responses: + '200': + description: Pet created + /pets/{petId}: + get: + summary: Info for a specific pet + responses: + '200': + description: Expected response to a valid request + """)) + def test_get_servers(self, mock_yaml_load, mock_open_file): + parser = OpenAPISpecificationParser(self.filepath) + servers = parser.get_servers() + self.assertEqual(servers, ["https://api.example.com", "https://staging.api.example.com"]) + + @patch("builtins.open", new_callable=mock_open, read_data="") + @patch("yaml.safe_load", return_value=yaml.safe_load(""" + openapi: 3.0.0 + info: + title: Sample API + version: 1.0.0 + servers: + - url: https://api.example.com + - url: https://staging.api.example.com + paths: + /pets: + get: + summary: List all pets + responses: + '200': + description: A paged array of pets + post: + summary: Create a pet + responses: + '200': + description: Pet created + /pets/{petId}: + get: + summary: Info for a specific pet + responses: + '200': + description: Expected response to a valid request + """)) + def test_get_paths(self, mock_yaml_load, mock_open_file): + parser = OpenAPISpecificationParser(self.filepath) + paths = parser.get_paths() + expected_paths = { + "/pets": { + "get": { + "summary": "List all pets", + "responses": { + "200": { + "description": "A paged array of pets" + } + } + }, + "post": { + "summary": "Create a pet", + "responses": { + "200": { + "description": "Pet created" + } + } + } + }, + "/pets/{petId}": { + "get": { + "summary": "Info for a specific pet", + "responses": { + "200": { + "description": "Expected response to a valid request" + } + } + } + } + } + self.assertEqual(paths, expected_paths) + + @patch("builtins.open", new_callable=mock_open, read_data="") + @patch("yaml.safe_load", return_value=yaml.safe_load(""" + openapi: 3.0.0 + info: + title: Sample API + version: 1.0.0 + servers: + - url: https://api.example.com + - url: https://staging.api.example.com + paths: + /pets: + get: + summary: List all pets + responses: + '200': + description: A paged array of pets + post: + summary: Create a pet + responses: + '200': + description: Pet created + /pets/{petId}: + get: + summary: Info for a specific pet + responses: + '200': + description: Expected response to a valid request + """)) + def test_get_operations(self, mock_yaml_load, mock_open_file): + parser = OpenAPISpecificationParser(self.filepath) + operations = parser.get_operations("/pets") + expected_operations = { + "get": { + "summary": "List all pets", + "responses": { + "200": { + "description": "A paged array of pets" + } + } + }, + "post": { + "summary": "Create a pet", + "responses": { + "200": { + "description": "Pet created" + } + } + } + } + self.assertEqual(operations, expected_operations) + + @patch("builtins.open", new_callable=mock_open, read_data="") + @patch("yaml.safe_load", return_value=yaml.safe_load(""" + openapi: 3.0.0 + info: + title: Sample API + version: 1.0.0 + servers: + - url: https://api.example.com + - url: https://staging.api.example.com + paths: + /pets: + get: + summary: List all pets + responses: + '200': + description: A paged array of pets + post: + summary: Create a pet + responses: + '200': + description: Pet created + /pets/{petId}: + get: + summary: Info for a specific pet + responses: + '200': + description: Expected response to a valid request + """)) + def test_print_api_details(self, mock_yaml_load, mock_open_file): + parser = OpenAPISpecificationParser(self.filepath) + with patch('builtins.print') as mocked_print: + parser.print_api_details() + mocked_print.assert_any_call("API Title:", "Sample API") + mocked_print.assert_any_call("API Version:", "1.0.0") + mocked_print.assert_any_call("Servers:", ["https://api.example.com", "https://staging.api.example.com"]) + mocked_print.assert_any_call("\nAvailable Paths and Operations:") + +if __name__ == "__main__": + unittest.main() From e4ef23a1f39e6d9af86a287c738ce56017e99cb2 Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 15:40:19 +0200 Subject: [PATCH 7/9] optimizeded code --- .../web_api_testing/prompt_engineer.py | 106 ++++++++---------- .../simple_openapi_documentation.py | 15 +-- .../web_api_testing/utils/response_handler.py | 2 +- 3 files changed, 58 insertions(+), 65 deletions(-) diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py b/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py index cd1bb40..8615e54 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py @@ -103,13 +103,45 @@ def get_http_action_template(self, method): else: return ( f"Create HTTPRequests of type {method} considering only the object with id=1 for the endpoint and understand the responses. Ensure that they are correct requests.") - - + def get_initial_steps(self, common_steps): + return [ + "Identify all available endpoints via GET Requests. Exclude those in this list: {self.found_endpoints}", + "Note down the response structures, status codes, and headers for each endpoint.", + "For each endpoint, document the following details: URL, HTTP method, query parameters and path variables, expected request body structure for requests, response structure for successful and error responses." + ] + common_steps + + def get_phase_steps(self, phase, common_steps): + if phase != "DELETE": + return [ + f"Identify for all endpoints {self.found_endpoints} excluding {self.endpoint_found_methods[phase]} a valid HTTP method {phase} call.", + self.get_http_action_template(phase) + ] + common_steps + else: + return [ + "Check for all endpoints the DELETE method. Delete the first instance for all endpoints.", + self.get_http_action_template(phase) + ] + common_steps + + def get_endpoints_needing_help(self): + endpoints_needing_help = [] + endpoints_and_needed_methods = {} + http_methods_set = {"GET", "POST", "PUT", "DELETE"} + + for endpoint, methods in self.endpoint_methods.items(): + missing_methods = http_methods_set - set(methods) + if len(methods) < 4: + endpoints_needing_help.append(endpoint) + endpoints_and_needed_methods[endpoint] = list(missing_methods) + + if endpoints_needing_help: + first_endpoint = endpoints_needing_help[0] + needed_method = endpoints_and_needed_methods[first_endpoint][0] + return [ + f"For endpoint {first_endpoint} find this missing method: {needed_method}. If all the HTTP methods have already been found for an endpoint, then do not include this endpoint in your search."] + return [] def chain_of_thought(self, doc=False, hint=""): """ Generates a prompt using the chain-of-thought strategy. - If 'doc' is True, it follows a detailed documentation-oriented prompt strategy based on the round number. - If 'doc' is False, it provides general guidance for early round numbers and focuses on HTTP methods for later rounds. Args: doc (bool): Determines whether the documentation-oriented chain of thought should be used. @@ -126,60 +158,19 @@ def chain_of_thought(self, doc=False, hint=""): "Make the OpenAPI specification available to developers by incorporating it into your API documentation site and keep the documentation up to date with API changes." ] - http_methods = [ "PUT", "DELETE"] - http_phase = { - 5: http_methods[0], - 10: http_methods[1] - } - + http_methods = ["PUT", "DELETE"] + http_phase = {10: http_methods[0], 15: http_methods[1]} if doc: - if self.round < 5: - - chain_of_thought_steps = [ - f"Identify all available endpoints via GET Requests. Exclude those in this list: {self.found_endpoints}", f"Note down the response structures, status codes, and headers for each endpoint.", - f"For each endpoint, document the following details: URL, HTTP method, " - f"query parameters and path variables, expected request body structure for requests, response structure for successful and error responses." - ] + common_steps + if self.round <= 5: + chain_of_thought_steps = self.get_initial_steps(common_steps) + elif self.round <= 10: + phase = http_phase.get(min(filter(lambda x: self.round <= x, http_phase.keys()))) + chain_of_thought_steps = self.get_phase_steps(phase, common_steps) else: - if self.round <= 10: - phase = http_phase.get(min(filter(lambda x: self.round <= x, http_phase.keys()))) - print(f'phase:{phase}') - if phase != "DELETE": - chain_of_thought_steps = [ - f"Identify for all endpoints {self.found_endpoints} excluding {self.endpoint_found_methods[phase]} a valid HTTP method {phase} call.", - self.get_http_action_template(phase) - ] + common_steps - else: - chain_of_thought_steps = [ - f"Check for all endpoints the DELETE method. Delete the first instance for all endpoints. ", - self.get_http_action_template(phase) - ] + common_steps - else: - endpoints_needing_help = [] - endpoints_and_needed_methods = {} - - # Standard HTTP methods - http_methods = {"GET", "POST", "PUT", "DELETE"} - - for endpoint in self.endpoint_methods: - # Calculate the missing methods for the current endpoint - missing_methods = http_methods - set(self.endpoint_methods[endpoint]) - - if len(self.endpoint_methods[endpoint]) < 4: - endpoints_needing_help.append(endpoint) - # Add the missing methods to the dictionary - endpoints_and_needed_methods[endpoint] = list(missing_methods) - - print(f'endpoints_and_needed_methods: {endpoints_and_needed_methods}') - print(f'first endpoint in list: {endpoints_needing_help[0]}') - print(f'methods needed for first endpoint: {endpoints_and_needed_methods[endpoints_needing_help[0]][0]}') - - chain_of_thought_steps = [f"For enpoint {endpoints_needing_help[0]} find this missing method :{endpoints_and_needed_methods[endpoints_needing_help[0]][0]} " - f"If all the HTTP methods have already been found for an endpoint, then do not include this endpoint in your search. ",] - + chain_of_thought_steps = self.get_endpoints_needing_help() else: if self.round == 0: - chain_of_thought_steps = ["Let's think step by step."] # Zero shot prompt + chain_of_thought_steps = ["Let's think step by step."] elif self.round <= 20: focus_phases = ["endpoints", "HTTP method GET", "HTTP method POST and PUT", "HTTP method DELETE"] focus_phase = focus_phases[self.round // 5] @@ -187,9 +178,10 @@ def chain_of_thought(self, doc=False, hint=""): else: chain_of_thought_steps = ["Look for exploits."] - print(f'chain of thought steps: {chain_of_thought_steps}') - prompt = self.check_prompt(self.previous_prompt, - chain_of_thought_steps + [hint] if hint else chain_of_thought_steps) + if hint: + chain_of_thought_steps.append(hint) + + prompt = self.check_prompt(self.previous_prompt, chain_of_thought_steps) return prompt def token_count(self, text): diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py index c1e9c41..285bd34 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py @@ -74,7 +74,7 @@ def _setup_initial_prompt(self): response_handler=self.response_handler) - def all_http_methods_found(self): + def all_http_methods_found(self,turn): print(f'found endpoints:{self.documentation_handler.endpoint_methods.items()}') print(f'found endpoints values:{self.documentation_handler.endpoint_methods.values()}') @@ -83,17 +83,20 @@ def all_http_methods_found(self): print(f'found endpoints:{found_endpoints}') print(f'expected endpoints:{expected_endpoints}') print(f'correct? {found_endpoints== expected_endpoints}') - if found_endpoints== expected_endpoints or found_endpoints == expected_endpoints -1: + if found_endpoints > 0 and (found_endpoints== expected_endpoints) : return True else: + if turn == 20: + if found_endpoints > 0 and (found_endpoints == expected_endpoints): + return True return False def perform_round(self, turn: int): prompt = self.prompt_engineer.generate_prompt(doc=True) response, completion = self.llm_handler.call_llm(prompt) - return self._handle_response(completion, response) + return self._handle_response(completion, response, turn) - def _handle_response(self, completion, response): + def _handle_response(self, completion, response, turn): message = completion.choices[0].message tool_call_id = message.tool_calls[0].id command = pydantic_core.to_json(response).decode() @@ -106,7 +109,6 @@ def _handle_response(self, completion, response): result_str = self.response_handler.parse_http_status_line(result) self._prompt_history.append(tool_message(result_str, tool_call_id)) invalid_flags = ["recorded","Not a valid HTTP method", "404" ,"Client Error: Not Found"] - print(f'result_str:{result_str}') if not result_str in invalid_flags or any(item in result_str for item in invalid_flags): self.prompt_engineer.found_endpoints = self.documentation_handler.update_openapi_spec(response, result) self.documentation_handler.write_openapi_to_yaml() @@ -120,8 +122,7 @@ def _handle_response(self, completion, response): http_methods_dict[method].append(endpoint) self.prompt_engineer.endpoint_found_methods = http_methods_dict self.prompt_engineer.endpoint_methods = self.documentation_handler.endpoint_methods - print(f'SCHEMAS:{self.prompt_engineer.schemas}') - return self.all_http_methods_found() + return self.all_http_methods_found(turn) diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py b/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py index 4b2dcae..da87481 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/utils/response_handler.py @@ -50,7 +50,7 @@ def parse_http_status_line(self, status_line): """ if status_line == "Not a valid HTTP method": return status_line - + status_line = status_line.split('\r\n')[0] # Regular expression to match valid HTTP status lines match = re.match(r'^(HTTP/\d\.\d) (\d{3}) (.*)$', status_line) if match: From d013162438c31bca484c8612ff2126a910619e8c Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 15:41:38 +0200 Subject: [PATCH 8/9] adjusted tests --- tests/test_web_api_documentation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_web_api_documentation.py b/tests/test_web_api_documentation.py index 8ef076a..ce70be6 100644 --- a/tests/test_web_api_documentation.py +++ b/tests/test_web_api_documentation.py @@ -35,8 +35,8 @@ def test_initial_prompt(self): 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.assertTrue(self.agent.all_http_methods_found()) + self.agent.all_http_methods_found(1) + self.assertFalse(self.agent.all_http_methods_found(1)) @patch('time.perf_counter', side_effect=[1, 2]) # Mocking perf_counter for consistent timing def test_perform_round(self, mock_perf_counter): @@ -61,7 +61,7 @@ def test_perform_round(self, mock_perf_counter): result = self.agent.perform_round(1) # Assertions - self.assertTrue(result) + self.assertFalse(result) # Check if the LLM was called with the correct parameters mock_create_with_completion = self.agent.llm.instructor.chat.completions.create_with_completion From 88fcf70a4d040bad95c52412fff2a474831c8aaa Mon Sep 17 00:00:00 2001 From: Diana Strauss Date: Tue, 6 Aug 2024 15:48:31 +0200 Subject: [PATCH 9/9] fixed wrong import --- tests/test_openapi_converter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_openapi_converter.py b/tests/test_openapi_converter.py index 67e1741..43354aa 100644 --- a/tests/test_openapi_converter.py +++ b/tests/test_openapi_converter.py @@ -4,7 +4,7 @@ import yaml import json -from hackingBuddyGPT.usecases.utils.openapi_converter import OpenAPISpecificationConverter +from hackingBuddyGPT.usecases.web_api_testing.utils.openapi_converter import OpenAPISpecificationConverter class TestOpenAPISpecificationConverter(unittest.TestCase):