From 651f4d82134e1a3f14add25c10f5e6eb0f28002a Mon Sep 17 00:00:00 2001
From: wwakabobik <i.vereshchagin@gmail.com>
Date: Tue, 17 Oct 2023 00:58:20 +0200
Subject: [PATCH] 0.2 - fix testing

---
 __init__.py                                   |   3 +-
 __main__.py                                   |  19 ++-
 examples/image_generation/gpt_functions.py    | 131 ++++++++++++++
 examples/speak_and_hear/test_gpt.py           |   4 +-
 examples/test_generator/generator_test.py     |  95 +++++++++++
 examples/test_generator/gpt_functions.py      | 161 +++++-------------
 examples/test_generator/pom_case_generator.py |  10 +-
 examples/test_generator/pytest_runner.py      |  55 ++++--
 examples/test_generator/tests/conftest.py     |  25 ++-
 requirements.txt                              |   8 +-
 utils/page_retriever.py                       |  23 ++-
 11 files changed, 366 insertions(+), 168 deletions(-)
 create mode 100644 examples/image_generation/gpt_functions.py
 create mode 100644 examples/test_generator/generator_test.py

diff --git a/__init__.py b/__init__.py
index 451bb7f..7fa3b4d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -5,7 +5,7 @@
 Copyright (c) 2023. All rights reserved.
 
 Created: 25.08.2023
-Last Modified: 12.09.2023
+Last Modified: 16.09.2023
 
 Description:
 This file is init point for project-wide structure.
@@ -14,6 +14,7 @@
 # Engines
 from .openai_api.src.openai_api.chatgpt import ChatGPT  # pylint: disable=unused-import
 from .openai_api.src.openai_api.dalle import DALLE  # pylint: disable=unused-import
+from .leonardo_api import Leonardo, LeonardoAsync  # pylint: disable=unused-import
 
 # Utils
 from .utils.tts import CustomTTS  # pylint: disable=unused-import
diff --git a/__main__.py b/__main__.py
index b0d22f8..1f4e81e 100644
--- a/__main__.py
+++ b/__main__.py
@@ -5,20 +5,21 @@
 Copyright (c) 2023. All rights reserved.
 
 Created: 25.08.2023
-Last Modified: 12.09.2023
+Last Modified: 16.09.2023
 
 Description:
 This file is entry point for project-wide structure.
 """
 
 # Engines
-from openai_api.src.openai_api.chatgpt import ChatGPT  # pylint: disable=unused-import
-from openai_api.src.openai_api.dalle import DALLE  # pylint: disable=unused-import
+from .openai_api.src.openai_api.chatgpt import ChatGPT  # pylint: disable=unused-import
+from .openai_api.src.openai_api.dalle import DALLE  # pylint: disable=unused-import
+from .leonardo_api import Leonardo, LeonardoAsync  # pylint: disable=unused-import
 
 # Utils
-from utils.tts import CustomTTS  # pylint: disable=unused-import
-from utils.transcriptors import CustomTranscriptor  # pylint: disable=unused-import
-from utils.translators import CustomTranslator  # pylint: disable=unused-import
-from utils.audio_recorder import AudioRecorder, record_and_convert_audio  # pylint: disable=unused-import
-from utils.logger_config import setup_logger  # pylint: disable=unused-import
-from utils.other import is_heroku_environment  # pylint: disable=unused-import
+from .utils.tts import CustomTTS  # pylint: disable=unused-import
+from .utils.transcriptors import CustomTranscriptor  # pylint: disable=unused-import
+from .utils.translators import CustomTranslator  # pylint: disable=unused-import
+from .utils.audio_recorder import AudioRecorder, record_and_convert_audio  # pylint: disable=unused-import
+from .utils.logger_config import setup_logger  # pylint: disable=unused-import
+from .utils.other import is_heroku_environment  # pylint: disable=unused-import
diff --git a/examples/image_generation/gpt_functions.py b/examples/image_generation/gpt_functions.py
new file mode 100644
index 0000000..7d66131
--- /dev/null
+++ b/examples/image_generation/gpt_functions.py
@@ -0,0 +1,131 @@
+import requests
+from PIL import Image
+from io import BytesIO
+
+import json
+
+from creds import oai_token, oai_organization
+from openai_api.src.openai_api.dalle import DALLE
+from leonardo_api.leonardo_sync import Leonardo
+from page_retriever import PageRetriever
+from pytest_runner import run_tests
+
+
+doc_engine = PageRetriever('https://wwakabobik.github.io/')
+
+
+def get_weather(city, units):
+    base_url = "http://api.openweathermap.org/data/2.5/weather"
+    params = {
+        "q": city,
+        "appid": "93171b03384f92ee3c55873452a49c7c",
+        "units": units
+    }
+    response = requests.get(base_url, params=params)
+    data = response.json()
+    return data
+
+
+def get_current_weather(location, unit="metric"):
+    """Get the current weather in a given location"""
+    owm_info = get_weather(location, units=unit)
+    weather_info = {
+        "location": location,
+        "temperature": owm_info["main"]["temp"],
+        "unit": unit,
+        "forecast": owm_info["weather"][0]["description"],
+        "wind": owm_info["wind"]["speed"]
+    }
+    return json.dumps(weather_info)
+
+
+def draw_image_using_dalle(prompt):
+    dalle = DALLE(auth_token=oai_token, organization=oai_organization)
+    image = dalle.create_image_url(prompt)
+    url_dict = {'image_url': image[0]}
+    response = requests.get(image[0])
+    img = Image.open(BytesIO(response.content))
+    img.show()
+    return json.dumps(url_dict)
+
+
+def draw_image(prompt):
+    leonardo = Leonardo(auth_token='a0178171-c67f-4922-afb3-458f24ecef1a')
+    leonardo.get_user_info()
+    response = leonardo.post_generations(prompt=prompt, num_images=1, guidance_scale=5,
+                                         model_id='e316348f-7773-490e-adcd-46757c738eb7', width=1024, height=768)
+    response = leonardo.wait_for_image_generation(generation_id=response['sdGenerationJob']['generationId'])
+    url_dict = {'image_url': response[0]['url']}
+    response = requests.get(url_dict['image_url'])
+    img = Image.open(BytesIO(response.content))
+    img.show()
+    return json.dumps(url_dict)
+
+
+gpt_functions = [
+        {
+            "name": "draw_image",
+            "description": "Draws image using user prompt. Returns url of image.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "prompt": {
+                        "type": "string",
+                        "description": "Prompt, the description, what should be drawn and how",
+                    },
+                },
+                "required": ["prompt"],
+            },
+        },
+        {
+            "name": "get_current_weather",
+            "description": "Get the current weather in a given location",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                "required": ["location"],
+            },
+        },
+        {
+            "name": "get_page_code",
+            "description": "Get page code to generate locators and tests",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "url": {
+                        "type": "string",
+                        "description": "The URL of the page to get the code from"
+                    }
+                },
+                "required": []
+            }
+        },
+        {
+            "name": "get_tests_results",
+            "description": "Get the results of the tests",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "test_files": {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "description": "The list of test files to run"
+                    }
+                },
+                "required": []
+            }
+        }
+    ]
+
+gpt_functions_dict = {'get_current_weather': get_current_weather,
+                      'draw_image': draw_image,
+                      'get_page_code': doc_engine.get_body_without_scripts,
+                      'get_tests_results': run_tests('tests/test_example.py')}
\ No newline at end of file
diff --git a/examples/speak_and_hear/test_gpt.py b/examples/speak_and_hear/test_gpt.py
index 83356e8..a5cd48d 100644
--- a/examples/speak_and_hear/test_gpt.py
+++ b/examples/speak_and_hear/test_gpt.py
@@ -20,8 +20,8 @@
 from utils.transcriptors import CustomTranscriptor
 from utils.tts import CustomTTS
 
-from creds import oai_token, oai_organization
-from openai_api.src.openai_api.chatgpt import ChatGPT
+from ..creds import oai_token, oai_organization
+from ...openai_api import ChatGPT
 
 
 gpt = ChatGPT(auth_token=oai_token, organization=oai_organization, model="gpt-3.5-turbo")
diff --git a/examples/test_generator/generator_test.py b/examples/test_generator/generator_test.py
new file mode 100644
index 0000000..d308896
--- /dev/null
+++ b/examples/test_generator/generator_test.py
@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+"""
+Filename: __generator_test__.py
+Author: Iliya Vereshchagin
+Copyright (c) 2023. All rights reserved.
+
+Created: 16.10.2023
+Last Modified: 16.10.2023
+
+Description:
+This file contains testing procedures for ChatGPt experiments
+"""
+
+import json
+import logging
+
+import asyncio
+
+from openai_api.src.openai_api import ChatGPT
+from openai_api.src.openai_api.logger_config import setup_logger
+from examples.creds import oai_token, oai_organization
+from examples.test_generator.pom_case_generator import PomTestCaseGenerator
+from examples.test_generator.gpt_functions import gpt_functions, gpt_functions_dict
+
+generator = PomTestCaseGenerator(url='https://www.saucedemo.com/')
+#generator = PomTestCaseGenerator(url='https://automationintesting.com/selenium/testpage/')
+
+
+system_instructions = """
+You're bot responsible for QA automation testing. You tech stack is selenium + pytest. I will provide you url for testing.
+
+1) You may obtain page code by calling "get_page_code" function. It will return you:
+ raw HTML document, what needs to be tested (guarded by ```). And you need to respond with json in following format:
+{
+"page_objects": [
+"@property\\n
+    def calculate_button(self):\\n
+        return WebDriverWait(self.driver, 10).until(\\n
+            EC.presence_of_element_located((By.XPATH, '//button[.='''Calculate''']'))\\n
+        )", <...>
+],
+"tests": ["def test_division_by_zero(page):\\n
+    page.numbers_input.send_keys(1024)\\n
+    page.divide_button.click()\\n
+    page.calculator_input.send_keys('0')\\n
+    page/calculate_button.click()\\n
+    assert page.error.text() == 'Error: divide by zero'", <...>],
+}
+This means you need to create page objects for each object on the page using laconic and stable XPATH locators (as short and stables as you can, use only By.XPATH locators, not By.ID, not By.CSS_SELECTOR or By.CLASS name), and then create all possible test cases for them. It might be some filed filling tests (errors, border checks, positive and negative cases), clicking, content changing, etc. Please respect to use 'page' fixture for every test, it's predefined in code and opens page under test before it.
+2) Then I may ask you to execute some tests. You can run demanded test via "get_tests_results" function, based on gathered content, you need to respond with json in following format:
+results = {
+    "passed": [],
+    "failed": [],
+    "error": [],
+    "failure details": {}
+}
+where "failure details" - is dict with keys equal to test names (which you generated) and possible failures details. If you got an failures and errors, you need to respond as in 1 with fixed code (page objects and/or tests).
+Answer only with JSON in format I mentioned in 1. Never add anything more than that (no explanations, no extra text, only json).
+3) In addition to 1 and 2 i may pass you extra info what kind of test data might be used (i.e. for form filling), but in general you need to generate all possible scenarios (valid/invalid/border cases, always add what's not listed by user, but should be for best quality of testing coverage).
+"""
+
+
+def setup_gpt():
+    """Setup GPT bot with appropriate functions and settings"""
+    gpt = ChatGPT(auth_token=oai_token, organization=oai_organization, model="gpt-4-0613")
+    gpt.logger = setup_logger("gpt", "gpt.log", logging.INFO)
+    gpt.system_settings = ""
+    gpt.function_dict = gpt_functions_dict
+    gpt.function_call = 'auto'
+    gpt.functions = gpt_functions
+    gpt.system_settings = system_instructions
+    return gpt
+
+
+async def main():
+    print("===Setup GPT bot===")
+    gpt = setup_gpt()
+    print("===Get page code of https://www.saucedemo.com/ and generate POM and tests===")
+    response = await anext(gpt.str_chat("Get page code of https://www.saucedemo.com/ and generate POM and tests"))
+    print(response)
+    response = response.replace('\n', '')
+    generator.create_files_from_json(json.loads(response),
+                                     pom_folder='examples/test_generator/pom',
+                                     tests_folder='examples/test_generator/tests')
+    print("===Get tests results for examples/test_generator/tests/test_index.py==")
+    response = await anext(gpt.str_chat("Get tests results for examples/test_generator/tests/test_index.py"))
+    print(response)
+    print("===If there are failures in code, please fix it by fixing POM and tests===")
+    response = await anext(gpt.str_chat("If there are failures in code, please fix it by fixing POM and tests"))
+    print(response)
+    generator.create_files_from_json(json.loads(response),
+                                     pom_folder='..pom',
+                                     tests_folder='examples/test_generator/tests')
+
+asyncio.run(main())
diff --git a/examples/test_generator/gpt_functions.py b/examples/test_generator/gpt_functions.py
index 7d66131..494a3bf 100644
--- a/examples/test_generator/gpt_functions.py
+++ b/examples/test_generator/gpt_functions.py
@@ -1,131 +1,48 @@
-import requests
-from PIL import Image
-from io import BytesIO
+# -*- coding: utf-8 -*-
+"""
+Filename: __gpt_functions__.py
+Author: Iliya Vereshchagin
+Copyright (c) 2023. All rights reserved.
 
-import json
+Created: 16.10.2023
+Last Modified: 16.10.2023
 
-from creds import oai_token, oai_organization
-from openai_api.src.openai_api.dalle import DALLE
-from leonardo_api.leonardo_sync import Leonardo
-from page_retriever import PageRetriever
-from pytest_runner import run_tests
-
-
-doc_engine = PageRetriever('https://wwakabobik.github.io/')
-
-
-def get_weather(city, units):
-    base_url = "http://api.openweathermap.org/data/2.5/weather"
-    params = {
-        "q": city,
-        "appid": "93171b03384f92ee3c55873452a49c7c",
-        "units": units
-    }
-    response = requests.get(base_url, params=params)
-    data = response.json()
-    return data
-
-
-def get_current_weather(location, unit="metric"):
-    """Get the current weather in a given location"""
-    owm_info = get_weather(location, units=unit)
-    weather_info = {
-        "location": location,
-        "temperature": owm_info["main"]["temp"],
-        "unit": unit,
-        "forecast": owm_info["weather"][0]["description"],
-        "wind": owm_info["wind"]["speed"]
-    }
-    return json.dumps(weather_info)
-
-
-def draw_image_using_dalle(prompt):
-    dalle = DALLE(auth_token=oai_token, organization=oai_organization)
-    image = dalle.create_image_url(prompt)
-    url_dict = {'image_url': image[0]}
-    response = requests.get(image[0])
-    img = Image.open(BytesIO(response.content))
-    img.show()
-    return json.dumps(url_dict)
-
-
-def draw_image(prompt):
-    leonardo = Leonardo(auth_token='a0178171-c67f-4922-afb3-458f24ecef1a')
-    leonardo.get_user_info()
-    response = leonardo.post_generations(prompt=prompt, num_images=1, guidance_scale=5,
-                                         model_id='e316348f-7773-490e-adcd-46757c738eb7', width=1024, height=768)
-    response = leonardo.wait_for_image_generation(generation_id=response['sdGenerationJob']['generationId'])
-    url_dict = {'image_url': response[0]['url']}
-    response = requests.get(url_dict['image_url'])
-    img = Image.open(BytesIO(response.content))
-    img.show()
-    return json.dumps(url_dict)
+Description:
+This file contains testing procedures for ChatGPt experiments
+"""
 
+from examples.test_generator.pytest_runner import run_tests
+from utils.page_retriever import PageRetriever
 
+doc_engine = PageRetriever("https://wwakabobik.github.io/")
 gpt_functions = [
-        {
-            "name": "draw_image",
-            "description": "Draws image using user prompt. Returns url of image.",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "prompt": {
-                        "type": "string",
-                        "description": "Prompt, the description, what should be drawn and how",
-                    },
-                },
-                "required": ["prompt"],
-            },
+    {
+        "name": "get_page_code",
+        "description": "Get page code to generate locators and tests",
+        "parameters": {
+            "type": "object",
+            "properties": {"url": {"type": "string", "description": "The URL of the page to get the code from"}},
+            "required": [],
         },
-        {
-            "name": "get_current_weather",
-            "description": "Get the current weather in a given location",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "location": {
-                        "type": "string",
-                        "description": "The city and state, e.g. San Francisco, CA",
-                    },
-                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
-                },
-                "required": ["location"],
+    },
+    {
+        "name": "get_tests_results",
+        "description": "Get the results of the tests",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "test_files": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "The list of test files to run",
+                }
             },
+            "required": [],
         },
-        {
-            "name": "get_page_code",
-            "description": "Get page code to generate locators and tests",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "url": {
-                        "type": "string",
-                        "description": "The URL of the page to get the code from"
-                    }
-                },
-                "required": []
-            }
-        },
-        {
-            "name": "get_tests_results",
-            "description": "Get the results of the tests",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "test_files": {
-                        "type": "array",
-                        "items": {
-                            "type": "string"
-                        },
-                        "description": "The list of test files to run"
-                    }
-                },
-                "required": []
-            }
-        }
-    ]
+    },
+]
 
-gpt_functions_dict = {'get_current_weather': get_current_weather,
-                      'draw_image': draw_image,
-                      'get_page_code': doc_engine.get_body_without_scripts,
-                      'get_tests_results': run_tests('tests/test_example.py')}
\ No newline at end of file
+gpt_functions_dict = {
+    "get_page_code": doc_engine.get_body_without_scripts,
+    "get_tests_results": run_tests,
+}
diff --git a/examples/test_generator/pom_case_generator.py b/examples/test_generator/pom_case_generator.py
index a66a08c..0561407 100644
--- a/examples/test_generator/pom_case_generator.py
+++ b/examples/test_generator/pom_case_generator.py
@@ -60,7 +60,7 @@ def ___create_test_file(file_name, tests, pom_folder='pom', tests_folder='tests'
             test_file.write('import pytest\n\n')
             test_file.write(f'from {pom_folder}.{os.path.splitext(f"page_{file_name}")[0]} import Page'
                             f'{"".join(word.capitalize() for word in file_name.split("_"))}\n\n\n')
-            test_file.write('@pytest.fixture(scope="module")\n')
+            test_file.write('@pytest.fixture(scope="function")\n')
             test_file.write('def page(driver):\n')
             test_file.write(f'    page_under_test = Page{"".join(word.capitalize() for word in file_name.split("_"))}(driver)\n')
             test_file.write(f'    driver.get(page_under_test.url)\n')
@@ -68,16 +68,18 @@ def ___create_test_file(file_name, tests, pom_folder='pom', tests_folder='tests'
             for test in tests:
                 test_file.write(f'{test}\n\n\n')
 
-    def create_files_from_json(self, json_data, url=''):
+    def create_files_from_json(self, json_data, url='', pom_folder='pom', tests_folder='tests'):
         """
         Create test and page object model files from json data.
 
         :param json_data: JSON data.
         :param url: URL of the page.
+        :param pom_folder: Folder for page object model files.
+        :param tests_folder: Folder for test files.
         """
         if not url:
             url = self.url
         parsed_url = urlparse(unquote(url))
         file_name = parsed_url.path.strip('/').replace('/', '_') or 'index'
-        self.___create_test_file(file_name, json_data['tests'], pom_folder='..pom')
-        self.___create_pom_file(file_name, json_data['page_objects'], url)
+        self.___create_test_file(file_name, json_data['tests'], pom_folder=f"..pom", tests_folder=tests_folder)
+        self.___create_pom_file(file_name, json_data['page_objects'], url, pom_folder=pom_folder)
diff --git a/examples/test_generator/pytest_runner.py b/examples/test_generator/pytest_runner.py
index f98c00b..a467ddf 100644
--- a/examples/test_generator/pytest_runner.py
+++ b/examples/test_generator/pytest_runner.py
@@ -1,21 +1,27 @@
 """This module runs pytest and returns the results in JSON format."""
+import io
+
 import json
 import pytest
-from pytest_jsonreport.plugin import JSONReport
+
+from utils.page_retriever import PageRetriever
 
 
-def run_tests(test_files):
+def run_tests(test_files, add_failed_html=True, add_failure_reasons=True, count_of_htmls=1):
     """
     Run tests and return results in JSON format.
 
     Args:
-        test_files: string with test files.
+        test_files: list with test files.
+        add_failed_html: boolean to add html report.
+        add_failure_reasons: boolean to add failure reasons.
+        count_of_htmls: count of htmls to add. Doesn't recommend to use more than 1.
 
     Returns:
         JSON with results.
 
     """
-    pytest.main(["-q", "--json-report", "--json-report-file=test_report.json"] + test_files)
+    pytest.main(["-q", "--json-report", "--json-report-file=test_report.json", "-n=4", "-rfEx --tb=none -p no:warnings -p no:logging"] + test_files)
 
     with open('test_report.json', encoding='utf-8') as json_file:
         data = json.load(json_file)
@@ -24,23 +30,40 @@ def run_tests(test_files):
         "passed": [],
         "failed": [],
         "error": [],
-        "failure details": {}
+        "failure details": {},
+        "failed_pages": {}
     }
 
     for test in data['tests']:
+        node_name = test['nodeid'].split('::')[1]
         if test['outcome'] == 'passed':
-            results["passed"].append(test['nodeid'])
-        elif test['outcome'] == 'failed':
-            results["failed"].append(test['nodeid'])
-            results["failure details"][test['nodeid']] = test['longrepr']
-            page_html = next((prop[1] for prop in test['user_properties'] if prop[0] == 'page_html'), None)
-            results["failed_pages"][test['nodeid']] = page_html
-        elif test['outcome'] == 'error':
-            results["error"].append(test['nodeid'])
-            results["failure details"][test['nodeid']] = test['longrepr']
-            page_html = next((prop[1] for prop in test['user_properties'] if prop[0] == 'page_html'), None)
-            results["failed_pages"][test['nodeid']] = page_html
+            results["passed"].append(node_name)
+        elif test['outcome'] == 'failed' or test['outcome'] == 'error':
+            results[test['outcome']].append(node_name)
+            if add_failure_reasons:
+                results["failure details"][node_name] = {node_name: test['call']['crash']}
+            if add_failed_html:
+                if len(results["failed_pages"]) < count_of_htmls:
+                    results["failed_pages"][node_name] = {node_name: parse_error_page(node_name)}
 
     json_results = json.dumps(results)
 
     return json_results
+
+
+def parse_error_page(node_name):
+    """
+    Parse error page.
+
+    Args:
+        node_name: name of the node.
+
+    Returns:
+        string with parsed page.
+    """
+    parser = PageRetriever()
+    try:
+        with open(f"{node_name}.html", "r", encoding="utf-8") as file:
+            return parser.remove_script_tags(parser.extract_body_content(file))
+    except io.UnsupportedOperation:
+        return "No page available."
diff --git a/examples/test_generator/tests/conftest.py b/examples/test_generator/tests/conftest.py
index 7d23610..a7e29b7 100644
--- a/examples/test_generator/tests/conftest.py
+++ b/examples/test_generator/tests/conftest.py
@@ -1,21 +1,30 @@
+from time import sleep
+
 import pytest
 
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.chrome.service import Service as ChromeService
 from webdriver_manager.chrome import ChromeDriverManager
 
 
+def pytest_runtest_makereport(item, call):
+    if "driver" in item.fixturenames:
+        web_driver = item.funcargs["driver"]
+        if call.when == "call" and call.excinfo is not None:
+            with open(f"{item.nodeid.split('::')[1]}.html", "w", encoding="utf-8") as file:
+                file.write(web_driver.page_source)
+
+
 @pytest.fixture
 def driver(request):
     options = Options()
     options.add_argument("--headless")
-    _driver = webdriver.Chrome(ChromeDriverManager().install(), options=options)
-
-    def save_page_source_on_failure():
-        if request.node.rep_call.failed or request.node.rep_call.error:
-            page_html = _driver.page_source
-            request.node.user_properties.append(("page_html", page_html))
+    options.headless = True
+    path = ChromeDriverManager().install()
+    _driver = webdriver.Chrome(service=ChromeService(executable_path=path, options=options), options=options)
 
-    request.addfinalizer(save_page_source_on_failure)
+    yield _driver
 
-    return _driver
+    _driver.close()
+    _driver.quit()
diff --git a/requirements.txt b/requirements.txt
index c099a52..7790654 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -20,4 +20,10 @@ readability==0.3.1
 # Testing
 webdriver_manager==4.0.1
 selenium==4.14.0
-pytest==7.4.2
\ No newline at end of file
+pytest==7.4.2
+pytest-json-report==1.5.0
+pytest-xdist=3.3.1
+# Third-party-test
+cohere==4.27
+llamaapi==0.1.36
+
diff --git a/utils/page_retriever.py b/utils/page_retriever.py
index 295b88a..f80baee 100644
--- a/utils/page_retriever.py
+++ b/utils/page_retriever.py
@@ -8,6 +8,7 @@
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.common.by import By
+from selenium.webdriver.chrome.service import Service as ChromeService
 from webdriver_manager.chrome import ChromeDriverManager
 
 
@@ -21,7 +22,9 @@ def __init__(self, url=''):
         """
         options = Options()
         options.add_argument("--headless")
-        self.driver = webdriver.Chrome(ChromeDriverManager().install(), options=options)
+        options.headless = True
+        path = ChromeDriverManager().install()
+        self.driver = webdriver.Chrome(service=ChromeService(executable_path=path), options=options)
         self.url = url
 
     def set_url(self, url):
@@ -32,28 +35,37 @@ def set_url(self, url):
         """
         self.url = url
 
-    def get_page(self):
+    def get_page(self, url=None):
         """
         Get the page content from the url.
 
+        :param url: URL of the page.
         :returns: HTML content of the page.
         """
+        if url:
+            self.set_url(url)
         return self.get_page_content(self.url)
 
-    def get_body(self):
+    def get_body(self, url=None):
         """
         Get the body content of the page.
 
+        :param url: URL of the page.
         :returns: Body content of the page.
         """
+        if url:
+            self.set_url(url)
         return self.extract_body_content(self.get_page())
 
-    def get_body_without_scripts(self):
+    def get_body_without_scripts(self, url=None):
         """
         Get the body content of the page without <script>...</script> tags.
 
+        :param url: URL of the page.
         :returns: Body content of the page without <script>...</script> tags.
         """
+        if url:
+            self.set_url(url)
         return self.remove_script_tags(self.get_body())
 
     def get_page_content(self, url):
@@ -73,10 +85,11 @@ def get_page_content(self, url):
                 "return window.performance.getEntriesByType('resource').filter(item => "
                 "item.initiatorType == 'xmlhttprequest' && item.duration == 0)"
             )
-            if not network_activity or time.time() - start_time > 30:  # Таймаут в 30 секунд
+            if not network_activity or time.time() - start_time > 30:
                 break
 
         content = self.driver.page_source
+        self.driver.close()
         self.driver.quit()
 
         return content