66
77import pytest
88from requests .sessions import Session
9- from requests_testadapter import TestAdapter , TestSession
9+ import requests_mock
1010
1111from invokeai .app .services .config import InvokeAIAppConfig
1212from invokeai .app .services .download import DownloadQueueService , DownloadQueueServiceBase
@@ -213,81 +213,83 @@ def mm2_model_manager(
213213@pytest .fixture
214214def mm2_session (embedding_file : Path , diffusers_dir : Path ) -> Session :
215215 """This fixtures defines a series of mock URLs for testing download and installation."""
216- sess : Session = TestSession ()
217- sess .mount (
216+ sess = Session ()
217+ adapter = requests_mock .Adapter ()
218+ sess .mount ("http://" , adapter )
219+ sess .mount ("https://" , adapter )
220+ adapter .register_uri (
221+ "GET" ,
218222 "https://test.com/missing_model.safetensors" ,
219- TestAdapter (
220- b"missing" ,
221- status = 404 ,
222- ),
223+ text = "missing" ,
224+ status_code = 404 ,
225+ reason = "NOT FOUND" ,
223226 )
224- sess .mount (
227+ adapter .register_uri (
228+ "GET" ,
225229 "https://huggingface.co/api/models/stabilityai/sdxl-turbo" ,
226- TestAdapter (
227- RepoHFMetadata1 ,
228- headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : len (RepoHFMetadata1 )},
229- ),
230- )
231- sess .mount (
232- "https://huggingface.co/api/models/stabilityai/sdxl-turbo-nofp16" ,
233- TestAdapter (
234- RepoHFMetadata1_nofp16 ,
235- headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : len (RepoHFMetadata1_nofp16 )},
236- ),
230+ content = RepoHFMetadata1 ,
231+ headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : str (len (RepoHFMetadata1 ))},
237232 )
238- sess .mount (
239- "https://civitai.com/api/v1/model-versions/242807" ,
240- TestAdapter (
241- RepoCivitaiVersionMetadata1 ,
233+ (
234+ adapter .register_uri (
235+ "GET" ,
236+ "https://huggingface.co/api/models/stabilityai/sdxl-turbo-nofp16" ,
237+ content = RepoHFMetadata1_nofp16 ,
242238 headers = {
243- "Content-Length" : len (RepoCivitaiVersionMetadata1 ),
239+ "Content-Type" : "application/json; charset=utf-8" ,
240+ "Content-Length" : str (len (RepoHFMetadata1_nofp16 )),
244241 },
245242 ),
246243 )
247- sess .mount (
244+ adapter .register_uri (
245+ "GET" ,
246+ "https://civitai.com/api/v1/model-versions/242807" ,
247+ content = RepoCivitaiVersionMetadata1 ,
248+ headers = {"Content-Length" : str (len (RepoCivitaiVersionMetadata1 ))},
249+ )
250+ adapter .register_uri (
251+ "GET" ,
248252 "https://civitai.com/api/v1/models/215485" ,
249- TestAdapter (
250- RepoCivitaiModelMetadata1 ,
251- headers = {
252- "Content-Length" : len (RepoCivitaiModelMetadata1 ),
253- },
254- ),
253+ content = RepoCivitaiModelMetadata1 ,
254+ headers = {"Content-Length" : str (len (RepoCivitaiModelMetadata1 ))},
255255 )
256- sess .mount (
256+ adapter .register_uri (
257+ "GET" ,
257258 "https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/model_index.json" ,
258- TestAdapter (
259- RepoHFModelJson1 ,
260- headers = {
261- "Content-Length" : len (RepoHFModelJson1 ),
262- },
263- ),
259+ content = RepoHFModelJson1 ,
260+ headers = {"Content-Length" : str (len (RepoHFModelJson1 ))},
264261 )
265262 with open (embedding_file , "rb" ) as f :
266263 data = f .read () # file is small - just 15K
267- sess .mount (
264+ adapter .register_uri (
265+ "GET" ,
268266 "https://www.test.foo/download/test_embedding.safetensors" ,
269- TestAdapter (data , headers = {"Content-Type" : "application/octet-stream" , "Content-Length" : len (data )}),
267+ content = data ,
268+ headers = {"Content-Type" : "application/octet-stream" , "Content-Length" : str (len (data ))},
270269 )
271- sess .mount (
270+ adapter .register_uri (
271+ "GET" ,
272272 "https://huggingface.co/api/models/stabilityai/sdxl-turbo" ,
273- TestAdapter (
274- RepoHFMetadata1 ,
275- headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : len (RepoHFMetadata1 )},
276- ),
273+ content = RepoHFMetadata1 ,
274+ headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : str (len (RepoHFMetadata1 ))},
275+ )
276+ adapter .register_uri (
277+ "GET" ,
278+ "https://huggingface.co/api/models/stabilityai/sdxl-turbo/revision/ModelRepoVariant.Default?blobs=True" ,
279+ content = RepoHFMetadata1 ,
280+ headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : str (len (RepoHFMetadata1 ))},
277281 )
278- sess .mount (
282+ adapter .register_uri (
283+ "GET" ,
279284 "https://huggingface.co/api/models/InvokeAI-test/textual_inversion_tests?blobs=True" ,
280- TestAdapter (
281- HFTestLoraMetadata ,
282- headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : len (HFTestLoraMetadata )},
283- ),
285+ content = HFTestLoraMetadata ,
286+ headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : str (len (HFTestLoraMetadata ))},
284287 )
285- sess .mount (
288+ adapter .register_uri (
289+ "GET" ,
286290 "https://huggingface.co/InvokeAI-test/textual_inversion_tests/resolve/main/learned_embeds-steps-1000.safetensors" ,
287- TestAdapter (
288- data ,
289- headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : len (data )},
290- ),
291+ content = data ,
292+ headers = {"Content-Type" : "application/json; charset=utf-8" , "Content-Length" : str (len (data ))},
291293 )
292294 for root , _ , files in os .walk (diffusers_dir ):
293295 for name in files :
@@ -296,55 +298,57 @@ def mm2_session(embedding_file: Path, diffusers_dir: Path) -> Session:
296298 url = f"https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/{ url_base } "
297299 with open (path , "rb" ) as f :
298300 data = f .read ()
299- sess .mount (
301+ adapter .register_uri (
302+ "GET" ,
300303 url ,
301- TestAdapter (
302- data ,
303- headers = {
304- "Content-Type" : "application/json; charset=utf-8" ,
305- "Content-Length" : len (data ),
306- },
307- ),
304+ content = data ,
305+ headers = {
306+ "Content-Type" : "application/json; charset=utf-8" ,
307+ "Content-Length" : str (len (data )),
308+ },
308309 )
309310
310311 for i in ["12345" , "9999" , "54321" ]:
311312 content = (
312313 b"I am a safetensors file " + bytearray (i , "utf-8" ) + bytearray (32_000 )
313314 ) # for pause tests, must make content large
314- sess .mount (
315+ adapter .register_uri (
316+ "GET" ,
315317 f"http://www.civitai.com/models/{ i } " ,
316- TestAdapter (
317- content ,
318- headers = {
319- "Content-Length" : len (content ),
320- "Content-Disposition" : f'filename="mock{ i } .safetensors"' ,
321- },
322- ),
318+ content = content ,
319+ headers = {
320+ "Content-Length" : str (len (content )),
321+ "Content-Disposition" : f'filename="mock{ i } .safetensors"' ,
322+ },
323323 )
324324
325- sess .mount (
325+ adapter .register_uri (
326+ "GET" ,
326327 "http://www.huggingface.co/foo.txt" ,
327- TestAdapter (
328- content ,
329- headers = {
330- "Content-Length" : len (content ),
331- "Content-Disposition" : 'filename="foo.safetensors"' ,
332- },
333- ),
328+ content = content ,
329+ headers = {
330+ "Content-Length" : str (len (content )),
331+ "Content-Disposition" : 'filename="foo.safetensors"' ,
332+ },
334333 )
335334
336335 # here are some malformed URLs to test
337336 # missing the content length
338- sess .mount (
337+ adapter .register_uri (
338+ "GET" ,
339339 "http://www.civitai.com/models/missing" ,
340- TestAdapter (
341- b"Missing content length" ,
342- headers = {
343- "Content-Disposition" : 'filename="missing.txt"' ,
344- },
345- ),
340+ text = "Missing content length" ,
341+ headers = {
342+ "Content-Disposition" : 'filename="missing.txt"' ,
343+ },
346344 )
347345 # not found test
348- sess .mount ("http://www.civitai.com/models/broken" , TestAdapter (b"Not found" , status = 404 ))
346+ adapter .register_uri (
347+ "GET" ,
348+ "http://www.civitai.com/models/broken" ,
349+ text = "Not found" ,
350+ status_code = 404 ,
351+ reason = "NOT FOUND" ,
352+ )
349353
350354 return sess
0 commit comments