1+ from __future__ import annotations
2+
13import urllib .request
4+ from collections .abc import Callable
5+ from dataclasses import dataclass , field
26from io import BytesIO
7+ from typing import Any
38
49import numpy as np
510import pytest
1116 POET_SKELETON_ID ,
1217 POETBackend ,
1318)
14- from dlclivegui .services .inference .models .poet .weights import WeightsDownloadWorker
19+ from dlclivegui .services .inference .models .poet .weights import (
20+ WeightsDownloadWorker ,
21+ )
22+
23+ TEST_WEIGHTS_URL = "https://example.invalid/weights"
24+
25+
26+ class FakeResponse (BytesIO ):
27+ """In-memory response implementing the context-manager protocol."""
28+
29+ def __init__ (self , content : bytes ) -> None :
30+ super ().__init__ (content )
31+ self .length = len (content )
32+
33+ def __enter__ (self ) -> FakeResponse :
34+ return self
35+
36+ def __exit__ (
37+ self ,
38+ exc_type ,
39+ exc_value ,
40+ traceback ,
41+ ) -> None :
42+ self .close ()
43+
44+
45+ @dataclass
46+ class UrlopenStub :
47+ """Configurable replacement for urllib.request.urlopen."""
48+
49+ content : bytes = b""
50+ error : Exception | None = None
51+ before_error : Callable [[], None ] | None = None
52+ expected_timeout : float = 30
53+ calls : list [tuple [Any , float | None ]] = field (default_factory = list )
54+
55+ def __call__ (
56+ self ,
57+ request ,
58+ * ,
59+ timeout : float | None = None ,
60+ ) -> FakeResponse :
61+ self .calls .append ((request , timeout ))
62+
63+ assert timeout == self .expected_timeout
64+
65+ if self .before_error is not None :
66+ self .before_error ()
67+
68+ if self .error is not None :
69+ raise self .error
70+
71+ return FakeResponse (self .content )
72+
73+
74+ @pytest .fixture
75+ def urlopen_factory (
76+ monkeypatch ,
77+ ) -> Callable [..., UrlopenStub ]:
78+ """Install and return a configurable urlopen test stub."""
79+
80+ def install (
81+ * ,
82+ content : bytes = b"" ,
83+ error : Exception | None = None ,
84+ before_error : Callable [[], None ] | None = None ,
85+ expected_timeout : float = 30 ,
86+ ) -> UrlopenStub :
87+ if error is None and before_error is not None :
88+ raise ValueError ("before_error requires an error." )
89+
90+ stub = UrlopenStub (
91+ content = content ,
92+ error = error ,
93+ before_error = before_error ,
94+ expected_timeout = expected_timeout ,
95+ )
96+
97+ monkeypatch .setattr (
98+ urllib .request ,
99+ "urlopen" ,
100+ stub ,
101+ )
102+
103+ return stub
104+
105+ return install
15106
16107
17108def test_poet_backend_rejects_missing_checkpoint (
@@ -44,7 +135,10 @@ def test_make_pose_packet_contains_poet_metadata(
44135 checkpoint .touch ()
45136
46137 backend = POETBackend (str (checkpoint ))
47- pose = np .zeros ((2 , 17 , 3 ), dtype = np .float32 )
138+ pose = np .zeros (
139+ (2 , 17 , 3 ),
140+ dtype = np .float32 ,
141+ )
48142
49143 packet = backend .make_pose_packet (pose )
50144
@@ -72,6 +166,7 @@ def test_make_pose_packet_without_detections_keeps_metadata(
72166 assert packet .keypoints is None
73167 assert packet .individual_ids == []
74168 assert packet .keypoint_names == list (POET_KEYPOINT_NAMES )
169+ assert packet .skeleton_id == POET_SKELETON_ID
75170 assert packet .skeleton_edges == POET_SKELETON_EDGES
76171
77172
@@ -83,12 +178,11 @@ def test_download_worker_reuses_existing_file(
83178 destination .write_bytes (b"existing" )
84179
85180 worker = WeightsDownloadWorker (
86- "https://example.invalid/weights" ,
181+ TEST_WEIGHTS_URL ,
87182 destination ,
88183 )
89184
90185 progress : list [int ] = []
91-
92186 worker .progress .connect (progress .append )
93187
94188 with qtbot .waitSignal (
@@ -102,82 +196,86 @@ def test_download_worker_reuses_existing_file(
102196 assert destination .read_bytes () == b"existing"
103197
104198
105- class FakeResponse (BytesIO ):
106- def __init__ (self , content : bytes ) -> None :
107- super ().__init__ (content )
108- self .length = len (content )
109-
110- def __enter__ (self ):
111- return self
112-
113- def __exit__ (
114- self ,
115- exc_type ,
116- exc_value ,
117- traceback ,
118- ) -> None :
119- self .close ()
120-
121-
122199def test_download_worker_writes_destination (
123200 qtbot ,
124- monkeypatch ,
125201 tmp_path ,
202+ urlopen_factory ,
126203) -> None :
127204 destination = tmp_path / "weights.pth"
128205 content = b"model-data"
129206
130- monkeypatch .setattr (
131- urllib .request ,
132- "urlopen" ,
133- lambda _request : FakeResponse (content ),
207+ urlopen_stub = urlopen_factory (
208+ content = content ,
134209 )
135210
136211 worker = WeightsDownloadWorker (
137- "https://example.invalid/weights" ,
212+ TEST_WEIGHTS_URL ,
138213 destination ,
139214 )
140215
216+ errors : list [str ] = []
217+ worker .error .connect (errors .append )
218+
141219 with qtbot .waitSignal (
142220 worker .finished ,
143221 timeout = 1000 ,
144222 ) as blocker :
145223 worker .run ()
146224
225+ assert errors == []
147226 assert blocker .args == [str (destination )]
148227 assert destination .read_bytes () == content
149228 assert not destination .with_suffix (".pth.part" ).exists ()
150229
230+ assert len (urlopen_stub .calls ) == 1
231+ request , timeout = urlopen_stub .calls [0 ]
232+
233+ assert isinstance (
234+ request ,
235+ urllib .request .Request ,
236+ )
237+ assert request .full_url == TEST_WEIGHTS_URL
238+ assert timeout == 30
239+
151240
152241def test_download_worker_reports_failure_and_cleans_partial_file (
153242 qtbot ,
154- monkeypatch ,
155243 tmp_path ,
244+ urlopen_factory ,
156245) -> None :
157246 destination = tmp_path / "weights.pth"
158247 partial_path = destination .with_suffix (".pth.part" )
159248
160- def failing_urlopen (_request ):
161- partial_path .write_bytes (b"partial" )
162- raise OSError ("download failed" )
163-
164- monkeypatch .setattr (
165- urllib .request ,
166- "urlopen" ,
167- failing_urlopen ,
249+ urlopen_stub = urlopen_factory (
250+ error = OSError ("download failed" ),
251+ before_error = lambda : partial_path .write_bytes (b"partial" ),
168252 )
169253
170254 worker = WeightsDownloadWorker (
171- "https://example.invalid/weights" ,
255+ TEST_WEIGHTS_URL ,
172256 destination ,
173257 )
174258
259+ finished_paths : list [str ] = []
260+ worker .finished .connect (finished_paths .append )
261+
175262 with qtbot .waitSignal (
176263 worker .error ,
177264 timeout = 1000 ,
178265 ) as blocker :
179266 worker .run ()
180267
181268 assert blocker .args == ["download failed" ]
269+ assert finished_paths == []
182270 assert not partial_path .exists ()
183271 assert not destination .exists ()
272+
273+ assert len (urlopen_stub .calls ) == 1
274+ request , timeout = urlopen_stub .calls [0 ]
275+
276+ assert isinstance (
277+ request ,
278+ urllib .request .Request ,
279+ )
280+ assert request .full_url == TEST_WEIGHTS_URL
281+ assert timeout == 30
0 commit comments