Untitled

import time
import unittest
from selenium import webdriver

class BaseTest(unittest.TestCase):
    def setUp(self):
        """Set up the WebDriver instance and configure test environment."""
        self.driver = webdriver.Chrome()
        self.driver.implicitly_wait(5)
        self.driver.maximize_window()

    def tearDown(self):
        """Tear down the WebDriver instance and capture screenshots on failure or success."""
        # Capture screenshot based on test outcome
        if self._test_has_failed():
            screenshot_filename = f"screenshot_{time.strftime('%Y%m%d-%H%M%S')}_failure.png"
            self.driver.save_screenshot(screenshot_filename)
            print(f"Failure screenshot {screenshot_filename} saved.")
        else:
            screenshot_filename = f"screenshot_{time.strftime('%Y%m%d-%H%M%S')}_success.png"
            self.driver.save_screenshot(screenshot_filename)
            print(f"Success screenshot {screenshot_filename} saved.")

        # Quit the WebDriver
        self.driver.quit()

    def _test_has_failed(self):
        """Check if the current test has failed."""
        return any(self._outcome.results.get(test, (None, None))[1] for test in self._outcome.results)