update testing framework

run_checks.sh

@@ -0,0 +1,3 @@
+python -m flake8 --extend-ignore=E203,W503 --max-line-length=120 klondike/ tests/
+python -m black klondike/ tests/
+python -m isort --profile black klondike/ tests/

tests/test_bigquery_connector.py

@@ -26,7 +26,7 @@ def _build_mock_cursor(self, query_results=None):
  cursor.fetchmany.side_effect = [query_results, []]
 
  if query_results:
- cursor.description = query_results
+ cursor.description = [(key, None) for key in query_results[0]]
 
  # Create a mock that will play the role of the connection
  connection = mock.MagicMock()
@@ -36,24 +36,27 @@ def _build_mock_cursor(self, query_results=None):
  client = mock.MagicMock()
 
  bq = BigQueryConnector()
+ bq.connection = connection
  bq._client = client
 
  return bq
 
  def test_read_dataframe_from_bigquery(self):
- # sql = "select * from my_table"
- # tbl = pa.table(
- # {
- # "city": ["Brooklyn", "San Francisco", "Richmond"],
- # "state": ["New York", "California", "Virginia"],
- # }
- # )
-
- # bq = self._build_mock_cursor(query_results=tbl)
- # df = bq.read_dataframe_from_bigquery(sql=sql)
-
- # assert isinstance(df, pl.DataFrame)
- pass
+ "Tests read functionality for the `BigQueryConnector` object"
+
+ sql = "select * from my_table"
+ tbl = [
+ {"city": "BROOKLYN", "state": "NY"},
+ {"city": "SAN FRANCSICO", "state": "CA"},
+ {"city": "RICHMOND", "state": "VA"},
+ ]
+ query_results = pa.RecordBatch.from_pylist(tbl)
+
+ bq = self._build_mock_cursor(query_results=query_results)
+ df = bq.read_dataframe_from_bigquery(sql=sql)
+
+ assert isinstance(df, pl.DataFrame)
 
  def test_write_dataframe_to_bigquery(self):
+ "Tests write functionality for the `BigQueryConnector` object"
  pass