@@ -526,41 +526,48 @@ def test_token_endpoint(self, mock_token_endpoint):
526526 endpoint = server .token_endpoint ()
527527 assert endpoint == "https://example.snowflakecomputing.com/"
528528
529- @patch ("rsconnect.api.get_connection_parameters " )
530- def test_token_endpoint_with_account (self , mock_get_connection_parameters ):
529+ @patch ("rsconnect.api.get_parameters " )
530+ def test_token_endpoint_with_account (self , mock_get_parameters ):
531531 server = SPCSConnectServer ("https://spcs.example.com" , "example_connection" )
532- mock_get_connection_parameters .return_value = {"account" : "test_account" }
532+ mock_get_parameters .return_value = {"account" : "test_account" }
533533 endpoint = server .token_endpoint ()
534534 assert endpoint == "https://test_account.snowflakecomputing.com/"
535- mock_get_connection_parameters .assert_called_once_with ("example_connection" )
535+ mock_get_parameters .assert_called_once_with ("example_connection" )
536536
537- @patch ("rsconnect.api.get_connection_parameters " )
538- def test_token_endpoint_with_none_params (self , mock_get_connection_parameters ):
537+ @patch ("rsconnect.api.get_parameters " )
538+ def test_token_endpoint_with_none_params (self , mock_get_parameters ):
539539 server = SPCSConnectServer ("https://spcs.example.com" , "example_connection" )
540- mock_get_connection_parameters .return_value = None
540+ mock_get_parameters .return_value = None
541541 with pytest .raises (RSConnectException , match = "No Snowflake connection found." ):
542542 server .token_endpoint ()
543543
544- @patch ("rsconnect.api.get_connection_parameters " )
545- def test_fmt_payload (self , mock_get_connection_parameters ):
544+ @patch ("rsconnect.api.get_parameters " )
545+ def test_fmt_payload (self , mock_get_parameters ):
546546 server = SPCSConnectServer ("https://spcs.example.com" , "example_connection" )
547- mock_get_connection_parameters .return_value = {"account" : "test_account" , "role" : "test_role" }
547+ mock_get_parameters .return_value = {
548+ "account" : "test_account" ,
549+ "role" : "test_role" ,
550+ "authenticator" : "SNOWFLAKE_JWT" ,
551+ }
548552
549553 with patch ("rsconnect.api.generate_jwt" ) as mock_generate_jwt :
550554 mock_generate_jwt .return_value = "mocked_jwt"
551555 payload = server .fmt_payload ()
552556
553- assert "scope=session%3Arole%3Atest_role+spcs.example.com" in payload
554- assert "assertion=mocked_jwt" in payload
555- assert "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer" in payload
557+ assert (
558+ payload ["body" ]
559+ == "scope=session%3Arole%3Atest_role+spcs.example.com&assertion=mocked_jwt&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer" # noqa
560+ )
561+ assert payload ["headers" ] == {"Content-Type" : "application/x-www-form-urlencoded" }
562+ assert payload ["path" ] == "/oauth/token"
556563
557- mock_get_connection_parameters .assert_called_once_with ("example_connection" )
564+ mock_get_parameters .assert_called_once_with ("example_connection" )
558565 mock_generate_jwt .assert_called_once_with ("example_connection" )
559566
560- @patch ("rsconnect.api.get_connection_parameters " )
561- def test_fmt_payload_with_none_params (self , mock_get_connection_parameters ):
567+ @patch ("rsconnect.api.get_parameters " )
568+ def test_fmt_payload_with_none_params (self , mock_get_parameters ):
562569 server = SPCSConnectServer ("https://spcs.example.com" , "example_connection" )
563- mock_get_connection_parameters .return_value = None
570+ mock_get_parameters .return_value = None
564571 with pytest .raises (RSConnectException , match = "No Snowflake connection found." ):
565572 server .fmt_payload ()
566573
@@ -579,7 +586,11 @@ def test_exchange_token_success(self, mock_fmt_payload, mock_token_endpoint, moc
579586
580587 # Mock the token endpoint and payload
581588 mock_token_endpoint .return_value = "https://example.snowflakecomputing.com/"
582- mock_fmt_payload .return_value = "mocked_payload"
589+ mock_fmt_payload .return_value = {
590+ "body" : "mocked_payload_body" ,
591+ "headers" : {"Content-Type" : "application/x-www-form-urlencoded" },
592+ "path" : "/oauth/token" ,
593+ }
583594
584595 # Call the method
585596 result = server .exchange_token ()
@@ -589,9 +600,9 @@ def test_exchange_token_success(self, mock_fmt_payload, mock_token_endpoint, moc
589600 mock_http_server .assert_called_once_with (url = "https://example.snowflakecomputing.com/" )
590601 mock_server_instance .request .assert_called_once_with (
591602 method = "POST" ,
592- path = "/oauth/token" ,
593- body = "mocked_payload" ,
603+ body = "mocked_payload_body" ,
594604 headers = {"Content-Type" : "application/x-www-form-urlencoded" },
605+ path = "/oauth/token" ,
595606 )
596607
597608 @patch ("rsconnect.api.HTTPServer" )
@@ -610,7 +621,11 @@ def test_exchange_token_error_status(self, mock_fmt_payload, mock_token_endpoint
610621
611622 # Mock the token endpoint and payload
612623 mock_token_endpoint .return_value = "https://example.snowflakecomputing.com/"
613- mock_fmt_payload .return_value = "mocked_payload"
624+ mock_fmt_payload .return_value = {
625+ "body" : "mocked_payload_body" ,
626+ "headers" : {"Content-Type" : "application/x-www-form-urlencoded" },
627+ "path" : "/oauth/token" ,
628+ }
614629
615630 # Call the method and verify it raises the expected exception
616631 with pytest .raises (RSConnectException , match = "Failed to exchange Snowflake token" ):
@@ -631,7 +646,11 @@ def test_exchange_token_empty_response(self, mock_fmt_payload, mock_token_endpoi
631646
632647 # Mock the token endpoint and payload
633648 mock_token_endpoint .return_value = "https://example.snowflakecomputing.com/"
634- mock_fmt_payload .return_value = "mocked_payload"
649+ mock_fmt_payload .return_value = {
650+ "body" : "mocked_payload_body" ,
651+ "headers" : {"Content-Type" : "application/x-www-form-urlencoded" },
652+ "path" : "/oauth/token" ,
653+ }
635654
636655 # Call the method and verify it raises the expected exception
637656 with pytest .raises (
0 commit comments