Skip to content

Commit

Permalink
feat: async client test fix
Browse files Browse the repository at this point in the history
  • Loading branch information
simjak committed Nov 12, 2024
1 parent 4e6e512 commit 536f2d2
Showing 1 changed file with 18 additions and 12 deletions.
30 changes: 18 additions & 12 deletions aurelio_sdk/client_async.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import asyncio
import io
import logging
import os
import time
Expand Down Expand Up @@ -221,12 +222,6 @@ async def extract_file(

extract_response = None

# If file is a file-like object, reset its position before retries
if hasattr(file, 'seek')
initial_file_position = file.tell()
else:
initial_file_position = None

async with aiohttp.ClientSession(timeout=session_timeout) as session:
for attempt in range(1, retries + 1):
# Form data
Expand Down Expand Up @@ -259,10 +254,21 @@ async def extract_file(
)
else:
logger.debug("Uploading file bytes")
# Reset file-like object position
if initial_file_position is not None:
file.seek(initial_file_position)
data.add_field("file", file)
try:
if isinstance(file, (bytes, bytearray)):
data.add_field("file", file)
elif isinstance(file, io.IOBase):
# Handle file-like objects
pos = file.tell()
file.seek(0)
file_content = file.read()
data.add_field("file", file_content)
file.seek(pos)
except Exception as e:
raise ApiError(
message=f"Failed to read file contents: {str(e)}",
base_url=self.base_url,
) from e

try:
async with session.post(
Expand Down Expand Up @@ -299,7 +305,7 @@ async def extract_file(
status_code=response.status,
)
except ApiRateLimitError as e:
raise e
raise e from None
except asyncio.TimeoutError:
if attempt == retries:
raise ApiTimeoutError(
Expand Down Expand Up @@ -429,7 +435,7 @@ async def extract_url(
status_code=response.status,
)
except ApiRateLimitError as e:
raise e
raise e from None
except asyncio.TimeoutError:
if attempt == retries:
raise ApiTimeoutError(
Expand Down

0 comments on commit 536f2d2

Please sign in to comment.