Skip to content

Commit

Permalink
Addressed PR suggestion.
Browse files Browse the repository at this point in the history
  • Loading branch information
utkarsharma2 committed Mar 23, 2022
1 parent 7248960 commit d181021
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
6 changes: 3 additions & 3 deletions conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,11 @@ def sample_dag():

@pytest.fixture
def test_table(request, sql_server):
table_type = True
is_tmp_table = True
table_options = {}

if getattr(request, "param", None):
table_type = request.param.get("is_temp", True)
is_tmp_table = request.param.get("is_temp", True)
table_options = request.param.get("param", {})
sql_name, hook = sql_server

Expand All @@ -91,7 +91,7 @@ def test_table(request, sql_server):
raise ValueError("Unsupported Database")

params.update(table_options)
table = TempTable(**params) if table_type else Table(**params)
table = TempTable(**params) if is_tmp_table else Table(**params)

yield table

Expand Down
3 changes: 2 additions & 1 deletion tests/operators/test_agnostic_load_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import os
import pathlib
import unittest.mock
from typing import Union
from unittest import mock

import pandas as pd
Expand Down Expand Up @@ -44,7 +45,7 @@
CWD = pathlib.Path(__file__).parent


def get_dataframe_from_table(sql_name: str, test_table: Table, hook):
def get_dataframe_from_table(sql_name: str, test_table: Union[Table, TempTable], hook):
if sql_name == "bigquery":
client = bigquery.Client()
query_job = client.query(
Expand Down

0 comments on commit d181021

Please sign in to comment.