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."""
        # Ensure that we have access to the test outcome
        if hasattr(self, '_outcome'):
            result = self._outcome

            # Generate timestamp for screenshot filenames
            timestamp = time.strftime("%Y%m%d-%H%M%S")

            # Capture a screenshot if there are any failures or errors
            if any(self._outcome.errors) or any(self._outcome.failures):
                screenshot_filename = f"screenshot_{timestamp}_failure.png"
                self.driver.save_screenshot(screenshot_filename)
                print(f"Failure screenshot {screenshot_filename} saved.")
            else:
                screenshot_filename = f"screenshot_{timestamp}_success.png"
                self.driver.save_screenshot(screenshot_filename)
                print(f"Success screenshot {screenshot_filename} saved.")

        # Quit the WebDriver
        self.driver.quit()