import re

def tab_code(txt):
    txt = txt.replace('  ', '    ')
    return txt


def mod_code(txt):
    new_text = ''
    for line in txt.splitlines():
        if 'def test_steam(self):' in line:
            new_text += '    @retry(tries=3, delay=10)' + '\n'
        if 'self.driver.find_element' in line:
              object_type = re.search(r'By\.\w*', line)[0]
              object_name = re.search(r'(, \"#?)([\w|\s]*)', line)[2]
              new_text += '        ' + f'WebDriverWait(self.driver, 30).until(EC.presence_of_element_located(({object_type}, {object_name}))' + '\n'
        new_text += line + '\n'
        if line == 'import json':
            new_text += 'import os' + '\n'
            new_text += 'from retry import retry' + '\n'
        if 'def setup_method(self, method):' in line:
            new_text += '        self.path = os.path.join(executable_path=r"C:\\Program Files (x86)\\Webdrivers\\chromedriver.exe")' + '\n'
    new_text = new_text.replace('self.driver = webdriver.Chrome()', 'self.driver = webdriver.Chrome(executable_path=self.path)')
    new_text = new_text.replace('from selenium.webdriver.support import expected_conditions',
                      'from selenium.webdriver.support import expected_conditions as EC')
    return new_text