From 66d79ebd1af7b79b70bca7ce4bb07cc718dd49ed Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Fri, 7 Jul 2023 15:59:33 +0530 Subject: [PATCH] Retain session.query in test --- tests/cli/commands/test_connection_command.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py index 7cc19af529d2a..0c48cd099b721 100644 --- a/tests/cli/commands/test_connection_command.py +++ b/tests/cli/commands/test_connection_command.py @@ -26,7 +26,6 @@ from unittest import mock import pytest -from sqlalchemy import select from airflow.cli import cli_parser from airflow.cli.commands import connection_command @@ -565,7 +564,7 @@ def test_cli_connection_add(self, cmd, expected_output, expected_conn): "port", "schema", ] - current_conn = session.scalars(select(Connection).where(Connection.conn_id == conn_id)).first() + current_conn = session.query(Connection).filter(Connection.conn_id == conn_id).first() assert expected_conn == {attr: getattr(current_conn, attr) for attr in comparable_attrs} def test_cli_connections_add_duplicate(self): @@ -675,7 +674,7 @@ def test_cli_delete_connections(self, session=None): assert "Successfully deleted connection with `conn_id`=new1" in stdout # Check deletions - result = session.scalars(select(Connection).filter(Connection.conn_id == "new1")).first() + result = session.query(Connection).filter(Connection.conn_id == "new1").first() assert result is None @@ -751,7 +750,7 @@ def test_cli_connections_import_should_load_connections(self, mock_exists, mock_ # Verify that the imported connections match the expected, sample connections with create_session() as session: - current_conns = session.scalars(select(Connection)).all() + current_conns = session.query(Connection).all() comparable_attrs = [ "conn_id", @@ -828,7 +827,7 @@ def test_cli_connections_import_should_not_overwrite_existing_connections( assert "Could not import connection new3: connection already exists." in stdout.getvalue() # Verify that the imported connections match the expected, sample connections - current_conns = session.scalars(select(Connection)).all() + current_conns = session.query(Connection).all() comparable_attrs = [ "conn_id", @@ -908,7 +907,7 @@ def test_cli_connections_import_should_overwrite_existing_connections( assert "Could not import connection new3: connection already exists." not in stdout.getvalue() # Verify that the imported connections match the expected, sample connections - current_conns = session.scalars(select(Connection)).all() + current_conns = session.query(Connection).all() comparable_attrs = [ "conn_id",