Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ pub fn execute(
if let Some(remaining) = remaining {
request = request.timeout(remaining);
}
let started = Instant::now();
let mut response = match request.send() {
Ok(response) => response,
Err(error) if error.is_timeout() => return Err("request timeout".into()),
Expand All @@ -88,7 +89,21 @@ pub fn execute(
.unwrap_or("")
.to_owned();
let mut bytes = Vec::new();
std::io::Read::take(&mut response, (512 << 20) + 1).read_to_end(&mut bytes)?;
if let Err(error) = std::io::Read::take(&mut response, (512 << 20) + 1).read_to_end(&mut bytes)
{
// The blocking response wraps body timeouts in io::Error; unlike send(),
// its Display text does not reliably include the underlying timeout.
if error.kind() == std::io::ErrorKind::TimedOut
|| error
.get_ref()
.and_then(|source| source.downcast_ref::<reqwest::Error>())
.is_some_and(reqwest::Error::is_timeout)
|| remaining.is_some_and(|budget| started.elapsed() >= budget)
{
return Err("request timeout".into());
}
return Err(error.into());
}
if bytes.len() > 512 << 20 {
return Err("API response exceeds 512 MiB".into());
}
Expand Down
15 changes: 13 additions & 2 deletions tests/cli_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import subprocess
import tempfile
import threading
import time
import unittest
from urllib.parse import parse_qs, urlsplit

Expand All @@ -18,7 +19,7 @@


@contextmanager
def api(responder=None):
def api(responder=None, body_delay=0):
calls=[]
class Handler(BaseHTTPRequestHandler):
def log_message(self,*_args):
Expand All @@ -33,7 +34,11 @@ def respond(self):
self.send_header('Content-Type',media)
self.send_header('Content-Length',str(len(data)))
self.end_headers()
self.wfile.write(data)
time.sleep(body_delay)
try:
self.wfile.write(data)
except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
pass # A timed-out client is expected to close this connection.
do_GET=respond
do_POST=respond
do_PATCH=respond
Expand Down Expand Up @@ -150,6 +155,12 @@ def test_timeout_and_api_failures_return_nonzero(self):
result=self.request(url,'balance')
self.assertEqual(result.returncode,1)
self.assertIn('API error 500: fixture failed',result.stderr)
def test_timeout_after_response_headers_is_identified(self):
with api(lambda calls:(200,dict(job_id='job-1') if len(calls)==1 else dict(status='processing'),'application/json'), body_delay=.5) as (url,calls):
result=self.request(url,'generate-image-v2','--body-json','{}','--wait','--timeout','100ms')
self.assertNotEqual(result.returncode,0)
self.assertIn('timeout',result.stderr)
self.assertEqual(len(calls),2)
def test_token_precedence_flag_environment_then_dotenv(self):
(self.root/'.env').write_text('export PIXELLAB_API_KEY="file-token"\n')
with api() as (url,calls):
Expand Down
Loading