Untitled

import time
import unittest
from selenium import webdriver

class BaseTest(unittest.TestCase):
    def setUp(self):
        self.driver = webdriver.Chrome()
        self.driver.implicitly_wait(5)
        self.driver.maximize_window()

    def tearDown(self):
        """Tear down the WebDriver instance and capture screenshots based on test results."""
        result = getattr(self, '_result', None)
        if result:
            # Check if there are any failures or errors
            failures_or_errors = result.failures or result.errors
            if failures_or_errors:
                timestamp = time.strftime("%Y%m%d-%H%M%S")
                screenshot_filename = f"screenshot_{timestamp}_failure.png"
                self.driver.save_screenshot(screenshot_filename)
                print(f"Failed screenshot {screenshot_filename} saved.")
            else:
                timestamp = time.strftime("%Y%m%d-%H%M%S")
                screenshot_filename = f"screenshot_{timestamp}_success.png"
                self.driver.save_screenshot(screenshot_filename)
                print(f"Success screenshot {screenshot_filename} saved.")
        
        self.driver.quit()