@@ -893,3 +893,76 @@ def test_get_columns(self, sea_client, sea_session_id, mock_cursor):
893893 cursor = mock_cursor ,
894894 )
895895 assert "Catalog name is required for get_columns" in str (excinfo .value )
896+
897+ def test_get_chunk_link (self , sea_client , mock_http_client , sea_command_id ):
898+ """Test get_chunk_link method."""
899+ # Setup mock response
900+ mock_response = {
901+ "external_links" : [
902+ {
903+ "external_link" : "https://example.com/data/chunk0" ,
904+ "expiration" : "2025-07-03T05:51:18.118009" ,
905+ "row_count" : 100 ,
906+ "byte_count" : 1024 ,
907+ "row_offset" : 0 ,
908+ "chunk_index" : 0 ,
909+ "next_chunk_index" : 1 ,
910+ "http_headers" : {"Authorization" : "Bearer token123" },
911+ }
912+ ]
913+ }
914+ mock_http_client ._make_request .return_value = mock_response
915+
916+ # Call the method
917+ result = sea_client .get_chunk_link ("test-statement-123" , 0 )
918+
919+ # Verify the HTTP client was called correctly
920+ mock_http_client ._make_request .assert_called_once_with (
921+ method = "GET" ,
922+ path = sea_client .CHUNK_PATH_WITH_ID_AND_INDEX .format (
923+ "test-statement-123" , 0
924+ ),
925+ )
926+
927+ # Verify the result
928+ assert result .external_link == "https://example.com/data/chunk0"
929+ assert result .expiration == "2025-07-03T05:51:18.118009"
930+ assert result .row_count == 100
931+ assert result .byte_count == 1024
932+ assert result .row_offset == 0
933+ assert result .chunk_index == 0
934+ assert result .next_chunk_index == 1
935+ assert result .http_headers == {"Authorization" : "Bearer token123" }
936+
937+ def test_get_chunk_link_not_found (self , sea_client , mock_http_client ):
938+ """Test get_chunk_link when the requested chunk is not found."""
939+ # Setup mock response with no matching chunk
940+ mock_response = {
941+ "external_links" : [
942+ {
943+ "external_link" : "https://example.com/data/chunk1" ,
944+ "expiration" : "2025-07-03T05:51:18.118009" ,
945+ "row_count" : 100 ,
946+ "byte_count" : 1024 ,
947+ "row_offset" : 100 ,
948+ "chunk_index" : 1 , # Different chunk index
949+ "next_chunk_index" : 2 ,
950+ "http_headers" : {"Authorization" : "Bearer token123" },
951+ }
952+ ]
953+ }
954+ mock_http_client ._make_request .return_value = mock_response
955+
956+ # Call the method and expect an exception
957+ with pytest .raises (
958+ ServerOperationError , match = "No link found for chunk index 0"
959+ ):
960+ sea_client .get_chunk_link ("test-statement-123" , 0 )
961+
962+ # Verify the HTTP client was called correctly
963+ mock_http_client ._make_request .assert_called_once_with (
964+ method = "GET" ,
965+ path = sea_client .CHUNK_PATH_WITH_ID_AND_INDEX .format (
966+ "test-statement-123" , 0
967+ ),
968+ )
0 commit comments