import re


def MigrateImports(content):
  """Updates import statements from TestNG to JUnit."""
  content_new = re.sub('org.testng.annotations.Test', 'org.junit.Test', content)

  content_new = re.sub('org.testng.annotations.BeforeMethod;',
                       'org.junit.Before;', content_new)

  content_new = re.sub('org.testng.annotations.BeforeClass;',
                       'org.junit.BeforeClass;', content_new)

  content_new = re.sub(
      'import org.testng.annotations.DataProvider;',
      '''import com.tngtech.java.junit.dataprovider.DataProvider;
import com.tngtech.java.junit.dataprovider.DataProviderRunner;
import com.tngtech.java.junit.dataprovider.UseDataProvider;
import org.junit.runner.RunWith;''', content_new)

  # for remaining imports such as assertEquals
  content_new = re.sub('testng', 'junit', content_new)

  return content_new