Skip to content

Commit 93d1d43

Browse files
remove un-necessary (?) changes
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 8afa396 commit 93d1d43

File tree

1 file changed

+73
-0
lines changed

1 file changed

+73
-0
lines changed

tests/unit/test_sea_backend.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)