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."""
        # Generate timestamp for screenshot filenames
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        
        # Check if any errors or failures occurred
        if self._resultForDoCleanups.errors or self._resultForDoCleanups.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()