Skip to content

Commit

Permalink
add unit test for existing query tag
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-kjimenezmorales committed Feb 5, 2024
1 parent f43dbde commit ec70f88
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 6 deletions.
2 changes: 1 addition & 1 deletion src/snowflake/snowpark/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -344,7 +344,7 @@ def create(self) -> "Session":
session = self._create_internal(self._options.get("connection"))

if self._app_name:
app_name_tag = f"APPNAME={self._app_name}"
app_name_tag = f'APPNAME={self._app_name}'
session.append_query_tag(app_name_tag)

return session
Expand Down
42 changes: 37 additions & 5 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -420,8 +420,8 @@ def test_connection_expiry():
m.assert_called_once()


def test_session_builder_app_name():
session = Session(
def test_session_builder_app_name_no_existing_query_tag():
mocked_session = Session(
ServerConnection(
{"": ""},
mock.Mock(
Expand All @@ -433,16 +433,48 @@ def test_session_builder_app_name():
),
),
)

mocked_session._get_remote_query_tag = MagicMock(return_value=None)

builder = Session.builder

with mock.patch.object(
builder,
"_create_internal",
return_value=session) as m:

return_value=mocked_session) as m:
app_name = 'my_app_name'
assert builder.app_name(app_name) is builder
created_session = builder.getOrCreate()
m.assert_called_once()
assert created_session.query_tag == f"APPNAME={app_name}"
assert created_session.query_tag == f'APPNAME={app_name}'


def test_session_builder_app_name_existing_query_tag():
mocked_session = Session(
ServerConnection(
{"": ""},
mock.Mock(
spec=SnowflakeConnection,
_telemetry=mock.Mock(),
_session_parameters=mock.Mock(),
is_closed=mock.Mock(return_value=False),
expired=False,
),
),
)

existing_query_tag = 'tag'

mocked_session._get_remote_query_tag = MagicMock(return_value=existing_query_tag)

builder = Session.builder

with mock.patch.object(
builder,
"_create_internal",
return_value=mocked_session) as m:
app_name = 'my_app_name'
assert builder.app_name(app_name) is builder
created_session = builder.getOrCreate()
m.assert_called_once()
assert created_session.query_tag == f'tag,APPNAME={app_name}'

0 comments on commit ec70f88

Please sign in to comment.