Skip to content

Commit b4ad201

Browse files
Fix test_useragent_header leaking agent env vars
Clear all KNOWN_AGENTS env vars in test_useragent_header before connecting, mirroring the pattern in test_agent_detection.py. Without this, running the test inside an AI-agent shell (e.g. Claude Code, which sets CLAUDECODE=1) causes session.py to append " agent/<product>" to the User-Agent and the strict `in` assertion fails. CI runners don't set agent env vars so this passes there, but the test is environment- sensitive locally. Co-authored-by: Isaac
1 parent cbd6a88 commit b4ad201

1 file changed

Lines changed: 4 additions & 1 deletion

File tree

tests/unit/test_session.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
THandleIdentifier,
99
)
1010
from databricks.sql.backend.types import SessionId, BackendType
11+
from databricks.sql.common.agent import KNOWN_AGENTS
1112
from databricks.sql.session import Session
1213

1314
import databricks.sql
@@ -97,7 +98,9 @@ def test_tls_arg_passthrough(self, mock_client_class, mock_http_client):
9798
assert kwargs["_tls_client_cert_key_password"] == "key password"
9899

99100
@patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
100-
def test_useragent_header(self, mock_client_class):
101+
def test_useragent_header(self, mock_client_class, monkeypatch):
102+
for env_var, _ in KNOWN_AGENTS:
103+
monkeypatch.delenv(env_var, raising=False)
101104
databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
102105

103106
call_kwargs = mock_client_class.call_args[1]

0 commit comments

Comments
 (0)