class TestBQIntegration():
def setup_class(self):
self.dataset_name = "gfluent_bq_dataset"
self.project_id = os.environ.get("PROJECT_ID")
self.sql = """
SELECT
exchange,
symbol,
enabled
FROM
market_data.required_products
LIMIT 5
"""
self.table_name = f"{self.dataset_name}.testing_bq_table"
self.bq = BQ(project=self.project_id)
self.bq.delete_dataset(self.dataset_name)
self.bq.create_dataset(self.dataset_name, location="EU")
def teardown_class(self):
self.bq.delete_dataset(self.dataset_name)
def test_query(self):
rows = self.bq.sql(self.sql).query()
assert rows.total_rows == 5
def test_query_load(self):
# load the query result to table
row_count = (self.bq.table(self.table_name).sql(self.sql).query())
# should only 5 rows
assert row_count == 5
# load again with append
row_count = (self.bq.table(self.table_name).sql(
self.sql).mode("WRITE_APPEND").query())
rows = BQ(
self.project_id).sql(f"select * from {self.table_name}").query()
assert rows.total_rows == 10
def test_truncate(self):
self.bq.table(self.table_name).truncate()
rows = BQ(
self.project_id).sql(f"select * from {self.table_name}").query()
assert rows.total_rows == 0
def test_delete(self):
self.bq.delete()
assert self.bq.is_exist() is False
def test_drop(self):
row_count = (self.bq.table(self.table_name).sql(self.sql).query())
assert self.bq.is_exist() is True
self.bq.drop()
assert self.bq.is_exist() is False
def test_is_exists(self):
assert self.bq.is_exist() is False
row_count = (self.bq.table(self.table_name).sql(self.sql).query())
assert self.bq.is_exist() is True
assert row_count == 5
def test_create_table(self):
table = f"{self.dataset_name}.students"
schema = [
bigquery.SchemaField(name="name",
field_type="STRING",
mode="REQUIRED",
description="student name"),
bigquery.SchemaField(name="age",
field_type="INTEGER",
mode="REQUIRED",
description="student age"),
]
bq = BQ(project=self.project_id).table(table).schema(schema)
assert bq.is_exist() is False
bq.create()
assert bq.is_exist() is True
class TestSheetIntegration():
def setup_class(self):
self.dataset_name = "gfluent_sheet_dataset"
self.project_id = os.environ.get("PROJECT_ID")
self.bq = BQ(project=self.project_id)
self.bq.delete_dataset(self.dataset_name)
self.bq.create_dataset(self.dataset_name, location="EU")
def teardown_class(self):
self.bq.delete_dataset(self.dataset_name)
def test_sheet_load(self):
sheet = Sheet(os.getenv("GOOGLE_APPLICATION_CREDENTIALS")).sheet_id(
TEST_UID).worksheet("data!A1:B4")
self.bq.table(f"{self.dataset_name}.target_table")
sheet.bq(self.bq).load(location="EU")
rows = BQ(project=self.project_id).sql(
f"select * from {self.dataset_name}.target_table").query()
assert rows.total_rows == 3
def test_sheet_load_with_schema(self):
_sheet = "no-header"
_range = "A:B"
_tb = "no_header"
sheet = Sheet(os.getenv("GOOGLE_APPLICATION_CREDENTIALS")).sheet_id(
TEST_UID).worksheet(_sheet).range(_range).schema([
bigquery.SchemaField(name="name",
field_type="STRING",
mode="REQUIRED",
description="student name"),
bigquery.SchemaField(name="age",
field_type="STRING",
mode="REQUIRED",
description="student age as string"),
])
self.bq.table(f"{self.dataset_name}.{_tb}")
sheet.bq(self.bq).load(location="EU")
rows = BQ(project=self.project_id).sql(
f"select * from {self.dataset_name}.{_tb}").query()
assert rows.total_rows == 3
def test_sheet_load_exist_exception(self):
_tb = "duplicated_table_name"
sheet = Sheet(os.getenv("GOOGLE_APPLICATION_CREDENTIALS")).sheet_id(
TEST_UID).worksheet("data!A1:B4")
self.bq.table(f"{self.dataset_name}.{_tb}")
# create the table
table = f"{self.dataset_name}.{_tb}"
schema = [
bigquery.SchemaField(name="name",
field_type="STRING",
mode="REQUIRED",
description="student name"),
bigquery.SchemaField(name="age",
field_type="INTEGER",
mode="REQUIRED",
description="student age"),
]
bq = BQ(project=self.project_id).table(table).schema(schema).create()
with pytest.raises(Conflict):
sheet.bq(self.bq).load(location="EU")