def test_db_driver(self, mock_db_driver, mock_UrlResolver, mock_os,
mock_get_config, mock_google_auth_default,
mock_google_cloud_storage_Client):
mock_creds = Mock(name='creds')
mock_db = Mock(name='db')
mock_url_resolver = mock_UrlResolver.return_value
mock_scratch_s3_url = mock_creds.default_scratch_s3_url.return_value
mock_s3_temp_base_loc = mock_url_resolver.directory_url.return_value
session = Session(creds=mock_creds)
out = session.db_driver(mock_db)
self.assertEqual(out, mock_db_driver.return_value)
mock_url_resolver.directory_url.assert_called_with(mock_scratch_s3_url)
mock_db_driver.assert_called_with(
db=mock_db,
url_resolver=mock_url_resolver,
s3_temp_base_loc=mock_s3_temp_base_loc)
class BaseRecordsIntegrationTest(unittest.TestCase):
def setUp(self):
# Ensure we're not getting any DWIM behavior out of the CLI
# session:
os.environ['RECORDS_MOVER_SESSION_TYPE'] = 'itest'
self.resources_dir = os.path.dirname(
os.path.abspath(__file__)) + '/../../resources'
self.session = Session(session_type='env',
default_db_creds_name=None,
default_aws_creds_name=None)
self.engine = self.session.get_default_db_engine()
self.driver = self.session.db_driver(self.engine)
if self.engine.name == 'bigquery':
self.schema_name = 'bq_itest'
# avoid per-table rate limits
elif self.engine.name == 'mysql':
self.schema_name = 'mysqlitest'
else:
self.schema_name = 'public'
table_name_prefix = "itest_"
build_num = os.environ.get("CIRCLE_BUILD_NUM", "local")
current_epoch = int(time.time())
self.table_name = f"{table_name_prefix}{build_num}_{current_epoch}"
self.fixture = RecordsDatabaseFixture(self.engine,
schema_name=self.schema_name,
table_name=self.table_name)
self.fixture.tear_down()
purge_old_tables(self.schema_name, table_name_prefix)
logger.debug("Initialized class!")
self.meta = MetaData()
self.records = self.session.records
def tearDown(self):
self.session = None
self.fixture.tear_down()
def table(self, schema, table):
return Table(table,
self.meta,
schema=schema,
autoload=True,
autoload_with=self.engine)
def variant_has_header(self, variant):
return variant in ['csv', 'bigquery']
def resource_name(self, format_type, variant, hints):
if hints.get('header-row', self.variant_has_header(variant)):
return f"{format_type}-{variant}-with-header"
else:
return f"{format_type}-{variant}-no-header"
def has_scratch_s3_bucket(self):
return os.environ.get('SCRATCH_S3_URL') is not None
def has_scratch_gcs_bucket(self):
return os.environ.get('SCRATCH_GCS_URL') is not None
def has_pandas(self):
try:
import pandas # noqa
logger.info("Just imported pandas")
return True
except ModuleNotFoundError:
logger.info("Could not find pandas")
return False
def unload_column_to_string(self, column_name: str,
records_format: BaseRecordsFormat) -> str:
targets = self.records.targets
sources = self.records.sources
with tempfile.TemporaryDirectory() as directory_name:
source = sources.table(schema_name=self.schema_name,
table_name=self.table_name,
db_engine=self.engine)
directory_url = pathlib.Path(directory_name).as_uri() + '/'
target = targets.directory_from_url(output_url=directory_url,
records_format=records_format)
self.records.move(source, target)
directory_loc = self.session.directory_url(directory_url)
records_dir = RecordsDirectory(records_loc=directory_loc)
with tempfile.NamedTemporaryFile() as t:
output_url = pathlib.Path(t.name).as_uri()
output_loc = self.session.file_url(output_url)
records_dir.save_to_url(output_loc)
return output_loc.string_contents()
class BaseRecordsIntegrationTest(unittest.TestCase):
def setUp(self):
# Ensure we're not getting any DWIM behavior out of the CLI
# session:
os.environ['RECORDS_MOVER_SESSION_TYPE'] = 'itest'
self.resources_dir = os.path.dirname(
os.path.abspath(__file__)) + '/../../resources'
self.session = Session(session_type='env',
default_db_creds_name=None,
default_aws_creds_name=None)
self.engine = self.session.get_default_db_engine()
self.driver = self.session.db_driver(self.engine)
if self.engine.name == 'bigquery':
self.schema_name = 'bq_itest'
# avoid per-table rate limits
elif self.engine.name == 'mysql':
self.schema_name = 'mysqlitest'
else:
self.schema_name = 'public'
table_name_prefix = "itest_"
build_num = os.environ.get("CIRCLE_BUILD_NUM", "local")
current_epoch = int(time.time())
self.table_name = f"{table_name_prefix}{build_num}_{current_epoch}"
self.fixture = RecordsDatabaseFixture(self.engine,
schema_name=self.schema_name,
table_name=self.table_name)
self.fixture.tear_down()
purge_old_tables(self.schema_name, table_name_prefix)
logger.debug("Initialized class!")
self.meta = MetaData()
self.records = self.session.records
def tearDown(self):
self.session = None
self.fixture.tear_down()
def table(self, schema, table):
return Table(table,
self.meta,
schema=schema,
autoload=True,
autoload_with=self.engine)
def variant_has_header(self, variant):
return variant in ['csv', 'bigquery']
def resource_name(self, format_type, variant, hints):
if hints.get('header-row', self.variant_has_header(variant)):
return f"{format_type}-{variant}-with-header"
else:
return f"{format_type}-{variant}-no-header"
def has_scratch_bucket(self):
return os.environ.get('SCRATCH_S3_URL') is not None
def has_pandas(self):
try:
import pandas # noqa
logger.info("Just imported pandas")
return True
except ModuleNotFoundError:
logger.info("Could not find pandas")
return False
