1717# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
1818# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
1919
20+ import contextlib
21+ import random
22+ import traceback
23+
24+ import sqlalchemy
25+ from sqlalchemy .testing import config
2026from sqlalchemy .testing .plugin .pytestplugin import * # noqa
2127from sqlalchemy .testing .plugin .pytestplugin import (
2228 pytest_sessionstart as _pytest_sessionstart ,
29+ pytest_sessionfinish as _pytest_sessionfinish ,
2330)
2431
2532import google .cloud .bigquery .dbapi .connection
2633import pybigquery .sqlalchemy_bigquery
27- import sqlalchemy
28- import traceback
2934
3035pybigquery .sqlalchemy_bigquery .BigQueryDialect .preexecute_autoincrement_sequences = True
3136google .cloud .bigquery .dbapi .connection .Connection .rollback = lambda self : None
@@ -51,18 +56,16 @@ def visit_delete(self, delete_stmt, *args, **kw):
5156pybigquery .sqlalchemy_bigquery .BigQueryCompiler .visit_delete = visit_delete
5257
5358
54- # Clean up test schemas so we don't get spurious errors when the tests
55- # try to create tables that already exist.
5659def pytest_sessionstart (session ):
57- client = google .cloud .bigquery .Client ()
58- for schema in "test_schema" , "test_pybigquery_sqla" :
59- for table_item in client .list_tables (f"{ client .project } .{ schema } " ):
60- table_id = table_item .table_id
61- list (
62- client .query (
63- f"drop { 'view' if table_id .endswith ('_v' ) else 'table' } "
64- f" { schema } .{ table_id } "
65- ).result ()
66- )
67- client .close ()
60+ dataset_id = f"test_pybigquery_sqla{ random .randint (0 , 1 << 63 )} "
61+ session .config .option .dburi = [f"bigquery:///{ dataset_id } " ]
62+ with contextlib .closing (google .cloud .bigquery .Client ()) as client :
63+ client .create_dataset (dataset_id )
6864 _pytest_sessionstart (session )
65+
66+
67+ def pytest_sessionfinish (session ):
68+ dataset_id = config .db .dialect .dataset_id
69+ _pytest_sessionfinish (session )
70+ with contextlib .closing (google .cloud .bigquery .Client ()) as client :
71+ client .delete_dataset (dataset_id , delete_contents = True )
0 commit comments