1469 lines
62 KiB
Python
1469 lines
62 KiB
Python
import unittest
|
||
from unittest.mock import patch, MagicMock
|
||
from core.type_utils import (
|
||
looks_like_model_file,
|
||
normalize_model_image,
|
||
fetch_remote_file_size,
|
||
get_version_sort_key,
|
||
check_credential_http,
|
||
extract_response_file_size,
|
||
normalize_sha256,
|
||
unique_ordered_strings,
|
||
extract_sha256_from_metadata,
|
||
extract_file_sha256,
|
||
extract_trained_words,
|
||
fetch_remote_file_size_cached,
|
||
clear_remote_size_cache,
|
||
normalize_category_to_model_type,
|
||
check_credential_preconditions,
|
||
)
|
||
from core.path_utils import (
|
||
calculate_file_sha256,
|
||
extract_safetensors_header_sha256,
|
||
extract_safetensors_header_metadata,
|
||
find_metadata_sidecar_path,
|
||
get_model_resolver_sidecar_path,
|
||
infer_safetensors_base_model,
|
||
read_json_safe,
|
||
read_merged_model_metadata,
|
||
write_json_atomic,
|
||
)
|
||
|
||
class UnifiedHelpersTests(unittest.TestCase):
|
||
|
||
def test_extract_file_sha256_preserves_provider_precedence(self):
|
||
self.assertEqual(
|
||
"direct",
|
||
extract_file_sha256({
|
||
"sha256": "direct",
|
||
"hash": "legacy",
|
||
"hashes": {"SHA256": "uppercase", "sha256": "lowercase"},
|
||
}),
|
||
)
|
||
self.assertEqual(
|
||
"legacy",
|
||
extract_file_sha256({
|
||
"hash": "legacy",
|
||
"hashes": {"SHA256": "uppercase", "sha256": "lowercase"},
|
||
}),
|
||
)
|
||
self.assertEqual(
|
||
"uppercase",
|
||
extract_file_sha256({"hashes": {"SHA256": "uppercase", "sha256": "lowercase"}}),
|
||
)
|
||
self.assertEqual("lowercase", extract_file_sha256({"hashes": {"sha256": "lowercase"}}))
|
||
|
||
def test_extract_file_sha256_handles_missing_or_invalid_shapes(self):
|
||
self.assertEqual("", extract_file_sha256(None))
|
||
self.assertEqual("", extract_file_sha256({"hashes": []}))
|
||
self.assertEqual("fallback", extract_file_sha256({"sha256": "", "hash": "fallback"}))
|
||
self.assertEqual("", extract_file_sha256({"sha256": 123}))
|
||
|
||
def test_extract_sha256_from_metadata_uses_first_valid_source(self):
|
||
preferred = "a" * 64
|
||
fallback = "b" * 64
|
||
|
||
self.assertEqual(
|
||
preferred,
|
||
extract_sha256_from_metadata({
|
||
"sha256": preferred,
|
||
"hash": fallback,
|
||
"hashes": {"SHA256": fallback},
|
||
"file_info": {"sha256": fallback},
|
||
}),
|
||
)
|
||
self.assertEqual(
|
||
fallback,
|
||
extract_sha256_from_metadata({
|
||
"sha256": "invalid",
|
||
"hash": fallback,
|
||
}),
|
||
)
|
||
self.assertEqual(
|
||
preferred,
|
||
extract_sha256_from_metadata({
|
||
"sha256": "invalid",
|
||
"hashes": {"SHA256": preferred, "sha256": fallback},
|
||
}),
|
||
)
|
||
self.assertEqual(
|
||
fallback,
|
||
extract_sha256_from_metadata({
|
||
"file_info": {"sha256": "invalid", "hash": fallback},
|
||
}),
|
||
)
|
||
|
||
def test_looks_like_model_file_huggingface(self):
|
||
# Valid HuggingFace URLs
|
||
self.assertTrue(looks_like_model_file("https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"))
|
||
self.assertTrue(looks_like_model_file("https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth"))
|
||
self.assertTrue(looks_like_model_file("hf://user/repo/model.safetensors"))
|
||
|
||
# Invalid / Spaces HuggingFace URLs
|
||
self.assertFalse(looks_like_model_file("https://huggingface.co/spaces/stabilityai/stable-diffusion"))
|
||
self.assertFalse(looks_like_model_file("https://huggingface.co/runwayml/stable-diffusion-v1-5"))
|
||
|
||
def test_looks_like_model_file_civitai_and_civarchive(self):
|
||
# Valid download URLs
|
||
self.assertTrue(looks_like_model_file("https://civitai.com/api/download/models/12345"))
|
||
self.assertTrue(looks_like_model_file("https://civitai.red/api/download/models/999"))
|
||
self.assertTrue(looks_like_model_file("https://civarchive.com/api/download/models/777"))
|
||
|
||
def test_looks_like_model_file_extensions(self):
|
||
# Known model extensions
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/my_model.safetensors"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/my_model.ckpt"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/my_model.gguf"))
|
||
|
||
# Unknown/invalid extensions
|
||
self.assertFalse(looks_like_model_file("https://example.com/files/my_model.txt"))
|
||
self.assertFalse(looks_like_model_file("https://example.com/files/image.png"))
|
||
self.assertFalse(looks_like_model_file("not_a_url_string"))
|
||
|
||
def test_looks_like_model_file_expected_filename(self):
|
||
# Matching expected filename
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/special_name.bin", "special_name.bin"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/special_name", "special_name"))
|
||
|
||
def test_normalize_model_image_civitai_format(self):
|
||
raw_image = {
|
||
"url": "https://image.civitai.com/x/width=1200/12345.jpeg",
|
||
"id": 12345,
|
||
"meta": {
|
||
"prompt": "masterpiece, best quality, 1girl",
|
||
"negativePrompt": "low quality, worst quality",
|
||
"cfgScale": 7.5,
|
||
"seed": 424242,
|
||
"steps": 25,
|
||
"sampler": "Euler a",
|
||
}
|
||
}
|
||
normalized = normalize_model_image(raw_image)
|
||
self.assertEqual(normalized["url"], "https://image.civitai.com/x/width=1200/12345.jpeg")
|
||
self.assertEqual(normalized["civitaiUrl"], "https://civitai.com/images/12345")
|
||
self.assertEqual(normalized["positive"], "masterpiece, best quality, 1girl")
|
||
self.assertEqual(normalized["negative"], "low quality, worst quality")
|
||
self.assertEqual(normalized["cfg"], 7.5)
|
||
self.assertEqual(normalized["seed"], 424242)
|
||
self.assertEqual(normalized["steps"], 25)
|
||
self.assertEqual(normalized["sampler"], "Euler a")
|
||
|
||
def test_normalize_model_image_preserves_video_type(self):
|
||
normalized = normalize_model_image({
|
||
"url": "https://image.civitai.com/x/original=true/12345.mp4",
|
||
"type": "video",
|
||
})
|
||
|
||
self.assertEqual(normalized["type"], "video")
|
||
|
||
def test_local_civitai_metadata_preserves_video_preview_path(self):
|
||
from core.sources.civitai import _metadata_to_model_info
|
||
|
||
result = _metadata_to_model_info({
|
||
"preview_url": "E:/Models/video-model.mp4",
|
||
"images": [{
|
||
"url": "https://image.civitai.com/x/original=true/12345.mp4",
|
||
"type": "video",
|
||
}],
|
||
})
|
||
|
||
self.assertEqual(result["preview_url"], "E:/Models/video-model.mp4")
|
||
self.assertEqual(result["images"][0]["type"], "video")
|
||
|
||
def test_local_civitai_metadata_separates_model_and_version_descriptions(self):
|
||
from core.sources.civitai import _metadata_to_model_info
|
||
|
||
result = _metadata_to_model_info({
|
||
"modelDescription": "Model page description",
|
||
"version_description": "Version-specific release notes",
|
||
"civitai": {
|
||
"description": "Nested version notes",
|
||
"model": {"description": "Nested model description"},
|
||
},
|
||
})
|
||
|
||
self.assertEqual("Model page description", result["description"])
|
||
self.assertEqual("Model page description", result["model_description"])
|
||
self.assertEqual(
|
||
"Version-specific release notes",
|
||
result["version_description"],
|
||
)
|
||
|
||
def test_normalize_model_image_preserves_metadata_when_normalized_again(self):
|
||
workflow = '{"workflow":{"nodes":[{"id":1,"type":"KSampler"}]}}'
|
||
normalized_image = {
|
||
"url": "https://image.civitai.com/x/width=1200/12345.jpeg",
|
||
"seed": 424242,
|
||
"metadata": {
|
||
"comfy": workflow,
|
||
"prompt": "portrait",
|
||
},
|
||
}
|
||
|
||
normalized_again = normalize_model_image(normalized_image)
|
||
|
||
self.assertEqual(normalized_again["metadata"]["comfy"], workflow)
|
||
self.assertEqual(normalized_again["positive"], "portrait")
|
||
self.assertEqual(normalized_again["seed"], 424242)
|
||
|
||
def test_normalize_model_image_civarchive_format(self):
|
||
raw_image = {
|
||
"imageUrl": "https://civarchive.com/img/9999.png",
|
||
"postUrl": "https://civitai.com/posts/8888",
|
||
"seed": 9999,
|
||
"steps": 30,
|
||
"cfg": 8.0,
|
||
"sampler": "DPM++ 2M Karras",
|
||
"positive": "cool car",
|
||
"negative": "blurry",
|
||
}
|
||
normalized = normalize_model_image(raw_image)
|
||
self.assertEqual(normalized["url"], "https://civarchive.com/img/9999.png")
|
||
self.assertEqual(normalized["civitaiUrl"], "https://civitai.com/posts/8888")
|
||
self.assertEqual(normalized["positive"], "cool car")
|
||
self.assertEqual(normalized["negative"], "blurry")
|
||
self.assertEqual(normalized["seed"], 9999)
|
||
self.assertEqual(normalized["steps"], 30)
|
||
self.assertEqual(normalized["cfg"], 8.0)
|
||
self.assertEqual(normalized["sampler"], "DPM++ 2M Karras")
|
||
|
||
def test_normalize_model_image_regex_fallback(self):
|
||
raw_image = {
|
||
"url": "https://image.civitai.com/images/width=1200/987654.jpeg",
|
||
}
|
||
normalized = normalize_model_image(raw_image)
|
||
self.assertEqual(normalized["civitaiUrl"], "https://civitai.com/images/987654")
|
||
|
||
@patch("requests.head")
|
||
@patch("requests.get")
|
||
def test_fetch_remote_file_size_head_success(self, mock_get, mock_head):
|
||
from requests.structures import CaseInsensitiveDict
|
||
# Setup mock response for HEAD
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.headers = CaseInsensitiveDict({"Content-Length": "104857600"})
|
||
mock_head.return_value = mock_response
|
||
|
||
size = fetch_remote_file_size("https://example.com/model.safetensors")
|
||
self.assertEqual(size, 104857600)
|
||
mock_head.assert_called_once()
|
||
mock_get.assert_not_called()
|
||
|
||
@patch("requests.head")
|
||
@patch("requests.get")
|
||
def test_fetch_remote_file_size_get_fallback(self, mock_get, mock_head):
|
||
from requests.structures import CaseInsensitiveDict
|
||
# HEAD fails
|
||
mock_head_response = MagicMock()
|
||
mock_head_response.status_code = 405 # Method Not Allowed
|
||
mock_head.return_value = mock_head_response
|
||
|
||
# GET succeeds with Range
|
||
mock_get_response = MagicMock()
|
||
mock_get_response.status_code = 206 # Partial Content
|
||
mock_get_response.headers = CaseInsensitiveDict({
|
||
"Content-Range": "bytes 0-0/209715200",
|
||
"Content-Length": "1"
|
||
})
|
||
mock_get.return_value = mock_get_response
|
||
|
||
size = fetch_remote_file_size("https://example.com/model.safetensors")
|
||
self.assertEqual(size, 209715200)
|
||
mock_head.assert_called_once()
|
||
mock_get.assert_called_once()
|
||
|
||
@patch("requests.head")
|
||
@patch("requests.get")
|
||
def test_fetch_remote_file_size_exception_handling(self, mock_get, mock_head):
|
||
mock_head.side_effect = Exception("Connection error")
|
||
mock_get.side_effect = Exception("Connection error")
|
||
|
||
size = fetch_remote_file_size("https://example.com/model.safetensors")
|
||
self.assertIsNone(size)
|
||
|
||
@patch("requests.head")
|
||
@patch("requests.get")
|
||
def test_fetch_remote_file_size_get_200_ok_fallback(self, mock_get, mock_head):
|
||
from requests.structures import CaseInsensitiveDict
|
||
# HEAD fails
|
||
mock_head_response = MagicMock()
|
||
mock_head_response.status_code = 405
|
||
mock_head.return_value = mock_head_response
|
||
|
||
# GET succeeds with 200 OK (full file)
|
||
mock_get_response = MagicMock()
|
||
mock_get_response.status_code = 200
|
||
mock_get_response.headers = CaseInsensitiveDict({
|
||
"Content-Length": "524288000"
|
||
})
|
||
mock_get.return_value = mock_get_response
|
||
|
||
size = fetch_remote_file_size("https://example.com/model.safetensors")
|
||
self.assertEqual(size, 524288000)
|
||
mock_head.assert_called_once()
|
||
mock_get.assert_called_once()
|
||
|
||
@patch("requests.head")
|
||
@patch("requests.get")
|
||
def test_fetch_remote_file_size_redirect_handling(self, mock_get, mock_head):
|
||
from requests.structures import CaseInsensitiveDict
|
||
# HEAD returns a 302 Redirect
|
||
mock_head_response = MagicMock()
|
||
mock_head_response.status_code = 302
|
||
mock_head_response.headers = CaseInsensitiveDict({
|
||
"Location": "https://redirected.com/real_model.safetensors"
|
||
})
|
||
|
||
# Redirects are followed manually so every destination can be validated.
|
||
mock_head_response_redirected = MagicMock()
|
||
mock_head_response_redirected.status_code = 200
|
||
mock_head_response_redirected.headers = CaseInsensitiveDict({
|
||
"Content-Length": "1000"
|
||
})
|
||
mock_head.side_effect = [mock_head_response, mock_head_response_redirected]
|
||
|
||
with patch(
|
||
"core.network_utils.validate_public_http_url",
|
||
side_effect=lambda url: url,
|
||
):
|
||
size = fetch_remote_file_size("https://example.com/redirect")
|
||
self.assertEqual(size, 1000)
|
||
|
||
self.assertEqual(mock_head.call_count, 2)
|
||
self.assertEqual(
|
||
mock_head.call_args_list[0].kwargs,
|
||
{
|
||
"headers": {"Accept-Encoding": "identity"},
|
||
"allow_redirects": False,
|
||
"stream": True,
|
||
"timeout": 15,
|
||
},
|
||
)
|
||
self.assertEqual(
|
||
mock_head.call_args_list[1].args[0],
|
||
"https://redirected.com/real_model.safetensors",
|
||
)
|
||
|
||
def test_write_json_atomic_recovery_on_replace_failure(self):
|
||
import tempfile
|
||
import os
|
||
from unittest.mock import patch
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "test.json")
|
||
with open(file_path, "w") as f:
|
||
f.write('{"initial": true}')
|
||
|
||
# Mock os.replace to raise OSError (simulating write/replace collision on Windows)
|
||
with patch("os.replace", side_effect=OSError("Permission Denied")):
|
||
with self.assertRaises(OSError):
|
||
write_json_atomic(file_path, {"updated": True})
|
||
|
||
# Verify the original file is untouched and readable
|
||
with open(file_path, "r") as f:
|
||
self.assertEqual(f.read(), '{"initial": true}')
|
||
|
||
# Verify no temporary files leak in the directory
|
||
self.assertEqual(os.listdir(tmpdir), ["test.json"])
|
||
|
||
def test_read_json_safe_malformed_syntax_fallback(self):
|
||
import tempfile
|
||
import os
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "corrupt.json")
|
||
with open(file_path, "w") as f:
|
||
f.write('{"unclosed_brace": true') # Malformed JSON
|
||
|
||
result = read_json_safe(file_path, {"fallback": True})
|
||
self.assertEqual(result, {"fallback": True})
|
||
|
||
def test_calculate_file_sha256_missing_file_handling(self):
|
||
# Empty path
|
||
self.assertIsNone(calculate_file_sha256(""))
|
||
# Non-existent path
|
||
self.assertIsNone(calculate_file_sha256("non_existent_file_path.safetensors"))
|
||
# A directory path
|
||
import tempfile
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
self.assertIsNone(calculate_file_sha256(tmpdir))
|
||
|
||
def test_get_version_sort_key_missing_and_invalid_timestamps(self):
|
||
# Normal sorting: timestamp, then id
|
||
v1 = {"id": 10, "published_at": "2024-01-01T00:00:00Z"}
|
||
v2 = {"id": 20, "published_at": "2024-01-01T00:00:00Z"}
|
||
# Both same timestamp, sorted by id
|
||
self.assertEqual(get_version_sort_key(v1), ("2024-01-01T00:00:00Z", 10))
|
||
self.assertEqual(get_version_sort_key(v2), ("2024-01-01T00:00:00Z", 20))
|
||
|
||
# Missing timestamp completely
|
||
v3 = {"id": 30}
|
||
self.assertEqual(get_version_sort_key(v3), ("", 30))
|
||
|
||
# None as timestamp
|
||
v4 = {"id": 40, "updatedAt": None}
|
||
self.assertEqual(get_version_sort_key(v4), ("", 40))
|
||
|
||
def test_get_version_sort_key_malformed_ids(self):
|
||
# String ID that cannot be cast to int should default to 0
|
||
v1 = {"id": "invalid_id_string", "published_at": "2024-01-01"}
|
||
self.assertEqual(get_version_sort_key(v1), ("2024-01-01", 0))
|
||
|
||
# Negative ID
|
||
v2 = {"id": -5, "published_at": "2024-01-01"}
|
||
self.assertEqual(get_version_sort_key(v2), ("2024-01-01", -5))
|
||
|
||
def test_looks_like_model_file_unexpected_edge_cases(self):
|
||
# URL with trailing text that invalidates the extension
|
||
self.assertFalse(looks_like_model_file("https://example.com/file.safetensors invalid"))
|
||
# No scheme
|
||
self.assertFalse(looks_like_model_file("example.com/file.safetensors"))
|
||
# No extension or path
|
||
self.assertFalse(looks_like_model_file("https://example.com/"))
|
||
# Empty string
|
||
self.assertFalse(looks_like_model_file(""))
|
||
|
||
def test_normalize_model_image_missing_civitai_urls(self):
|
||
# Image URL contains numbers at the end, should extract post ID
|
||
raw_image = {
|
||
"url": "https://image.civitai.com/x/width=1200/987654.jpeg",
|
||
}
|
||
normalized = normalize_model_image(raw_image)
|
||
self.assertEqual(normalized["civitaiUrl"], "https://civitai.com/images/987654")
|
||
|
||
# Non-CivitAI image URL with no numbers shouldn't build civitaiUrl
|
||
raw_image_other = {
|
||
"url": "https://example.com/some_image.png",
|
||
}
|
||
normalized_other = normalize_model_image(raw_image_other)
|
||
self.assertEqual(normalized_other["civitaiUrl"], "")
|
||
|
||
@patch("requests.post")
|
||
@patch("requests.get")
|
||
def test_check_credential_http_various_responses(self, mock_get, mock_post):
|
||
# 1. Success case with GET and username parsing
|
||
mock_response_200 = MagicMock()
|
||
mock_response_200.status_code = 200
|
||
mock_response_200.json.return_value = {"user": {"name": "civit_maker"}}
|
||
mock_get.return_value = mock_response_200
|
||
|
||
result = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={"Authorization": "Bearer key"},
|
||
get_username=lambda data: data["user"]["name"],
|
||
success_message="Authenticated"
|
||
)
|
||
self.assertTrue(result["success"])
|
||
self.assertTrue(result["valid"])
|
||
self.assertEqual(result["username"], "civit_maker")
|
||
self.assertEqual(result["message"], "Authenticated for civit_maker.")
|
||
|
||
# 2. HTTP 401/403 Invalid case
|
||
mock_response_401 = MagicMock()
|
||
mock_response_401.status_code = 401
|
||
mock_get.return_value = mock_response_401
|
||
|
||
result_401 = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={"Authorization": "Bearer bad"},
|
||
error_msg_401_403="Denied Access"
|
||
)
|
||
self.assertTrue(result_401["success"])
|
||
self.assertFalse(result_401["valid"])
|
||
self.assertEqual(result_401["status"], "invalid")
|
||
self.assertEqual(result_401["message"], "Denied Access")
|
||
self.assertEqual(result_401["status_code"], 401)
|
||
|
||
# 3. HTTP 429 Rate Limit case with custom handler
|
||
mock_response_429 = MagicMock()
|
||
mock_response_429.status_code = 429
|
||
mock_get.return_value = mock_response_429
|
||
|
||
custom_handler = lambda r: {"status": "limited", "retry": True}
|
||
result_429 = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={},
|
||
custom_429_handler=custom_handler
|
||
)
|
||
self.assertEqual(result_429, {"status": "limited", "retry": True})
|
||
|
||
# 4. Timeout Exception handling
|
||
import requests
|
||
mock_get.side_effect = requests.exceptions.Timeout("Request timed out")
|
||
result_timeout = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={}
|
||
)
|
||
self.assertFalse(result_timeout["success"])
|
||
self.assertEqual(result_timeout["status"], "timeout")
|
||
|
||
def test_calculate_file_sha256_success(self):
|
||
import tempfile
|
||
import os
|
||
import hashlib
|
||
|
||
content = b"model binary data block"
|
||
expected_hash = hashlib.sha256(content).hexdigest()
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "mock_model.safetensors")
|
||
with open(file_path, "wb") as f:
|
||
f.write(content)
|
||
|
||
calculated_hash = calculate_file_sha256(file_path, chunk_size=4)
|
||
self.assertEqual(calculated_hash, expected_hash)
|
||
|
||
def write_safetensors_stub(self, file_path, header):
|
||
import json
|
||
|
||
header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
||
with open(file_path, "wb") as f:
|
||
f.write(len(header_bytes).to_bytes(8, byteorder="little"))
|
||
f.write(header_bytes)
|
||
f.write(b"tensor payload")
|
||
|
||
def test_extract_safetensors_header_sha256(self):
|
||
import os
|
||
import tempfile
|
||
|
||
expected_hash = "a" * 64
|
||
header = {"__metadata__": {"modelspec.hash.sha256": expected_hash.upper()}}
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "header_hash.safetensors")
|
||
self.write_safetensors_stub(file_path, header)
|
||
|
||
self.assertEqual(extract_safetensors_header_sha256(file_path), expected_hash)
|
||
|
||
with patch(
|
||
"core.path_utils.hashlib.sha256",
|
||
side_effect=AssertionError("full file hash should not be calculated"),
|
||
):
|
||
self.assertEqual(calculate_file_sha256(file_path), expected_hash)
|
||
|
||
def test_extract_safetensors_header_sha256_with_0x_prefix(self):
|
||
import os
|
||
import tempfile
|
||
|
||
expected_hash = "b" * 64
|
||
header = {"__metadata__": {"modelspec.hash_sha256": "0x" + expected_hash}}
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "header_hash_0x.safetensors")
|
||
self.write_safetensors_stub(file_path, header)
|
||
|
||
self.assertEqual(extract_safetensors_header_sha256(file_path), expected_hash)
|
||
|
||
def test_extract_safetensors_header_sha256_ignores_sshs_model_hash(self):
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import tempfile
|
||
|
||
header = {
|
||
"__metadata__": {
|
||
"sshs_model_hash": "c" * 64,
|
||
"sshs_legacy_hash": "dacc4e08",
|
||
},
|
||
}
|
||
header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
||
payload = (
|
||
len(header_bytes).to_bytes(8, byteorder="little")
|
||
+ header_bytes
|
||
+ b"tensor payload"
|
||
)
|
||
expected_hash = hashlib.sha256(payload).hexdigest()
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "sshs_hash.safetensors")
|
||
with open(file_path, "wb") as f:
|
||
f.write(payload)
|
||
|
||
self.assertIsNone(extract_safetensors_header_sha256(file_path))
|
||
self.assertEqual(calculate_file_sha256(file_path), expected_hash)
|
||
|
||
def test_infer_safetensors_base_model_from_architecture(self):
|
||
import os
|
||
import tempfile
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "sdxl_arch.safetensors")
|
||
self.write_safetensors_stub(
|
||
file_path,
|
||
{
|
||
"__metadata__": {
|
||
"modelspec.architecture": "stable-diffusion-xl-v1-base",
|
||
},
|
||
},
|
||
)
|
||
|
||
self.assertEqual(infer_safetensors_base_model(file_path), "SDXL 1.0")
|
||
|
||
def test_extract_safetensors_header_metadata_summary(self):
|
||
import json
|
||
import os
|
||
import tempfile
|
||
|
||
expected_hash = "c" * 64
|
||
model_hash = "d" * 64
|
||
header = {
|
||
"__metadata__": {
|
||
"modelspec.title": "Header Model",
|
||
"modelspec.author": "Header Author",
|
||
"modelspec.description": "Header description",
|
||
"modelspec.license": "MIT",
|
||
"modelspec.hash.sha256": expected_hash,
|
||
"modelspec.tags": json.dumps(["style", "character"]),
|
||
"modelspec.trigger_phrase": "hero, dramatic light",
|
||
"modelspec.thumbnail": "data:image/png;base64,abc",
|
||
"ss_base_model_version": "pony",
|
||
"ss_clip_skip": "2",
|
||
"ss_network_module": "networks.lora",
|
||
"ss_network_dim": "16",
|
||
"sshs_model_hash": model_hash,
|
||
"ss_tag_frequency": json.dumps(
|
||
{"dataset": {"blue eyes": 12, "solo": 9}}
|
||
),
|
||
"workflow": "{}",
|
||
},
|
||
}
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "header_metadata.safetensors")
|
||
self.write_safetensors_stub(file_path, header)
|
||
|
||
metadata = extract_safetensors_header_metadata(file_path)
|
||
|
||
self.assertTrue(metadata["from_safetensors_header"])
|
||
self.assertEqual(metadata["model_name"], "Header Model")
|
||
self.assertEqual(metadata["creator"]["username"], "Header Author")
|
||
self.assertEqual(metadata["description"], "Header description")
|
||
self.assertEqual(metadata["license"], "MIT")
|
||
self.assertEqual(metadata["base_model"], "Pony")
|
||
self.assertEqual(metadata["sha256"], expected_hash)
|
||
self.assertEqual(metadata["sha256_source"], "safetensors_header")
|
||
self.assertEqual(metadata["clip_skip"], 2)
|
||
self.assertIn("style", metadata["tags"])
|
||
self.assertIn("blue eyes", metadata["tags"])
|
||
self.assertEqual(metadata["trained_words"], ["hero", "dramatic light"])
|
||
self.assertEqual(metadata["model_type"], "LORA")
|
||
self.assertTrue(metadata["metadata_summary"]["has_embedded_workflow"])
|
||
self.assertEqual(metadata["metadata_summary"]["network_dim"], "16")
|
||
self.assertEqual(metadata["metadata_summary"]["model_hash"], model_hash)
|
||
|
||
def test_infer_safetensors_base_model_from_tensor_fingerprints(self):
|
||
import os
|
||
import tempfile
|
||
|
||
cases = [
|
||
("flux.safetensors", "double_blocks.0.img_attn.qkv.weight", "Flux.1 D"),
|
||
("sd3.safetensors", "joint_blocks.0.x_block.attn.qkv.weight", "SD 3"),
|
||
("sdxl.safetensors", "conditioner.embedders.1.model.transformer", "SDXL 1.0"),
|
||
(
|
||
"sd15.safetensors",
|
||
"model.diffusion_model.input_blocks.0.0.weight",
|
||
"SD 1.5",
|
||
),
|
||
]
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
for filename, tensor_key, expected_base_model in cases:
|
||
file_path = os.path.join(tmpdir, filename)
|
||
self.write_safetensors_stub(
|
||
file_path,
|
||
{
|
||
"__metadata__": {},
|
||
tensor_key: {"dtype": "F16", "shape": [1]},
|
||
},
|
||
)
|
||
self.assertEqual(
|
||
infer_safetensors_base_model(file_path),
|
||
expected_base_model,
|
||
)
|
||
|
||
def test_calculate_file_sha256_falls_back_without_header_sha256(self):
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import tempfile
|
||
|
||
header = json.dumps(
|
||
{"__metadata__": {"modelspec.hash.blake3": "b" * 64}},
|
||
separators=(",", ":"),
|
||
).encode("utf-8")
|
||
payload = (
|
||
len(header).to_bytes(8, byteorder="little")
|
||
+ header
|
||
+ b"tensor payload"
|
||
)
|
||
expected_hash = hashlib.sha256(payload).hexdigest()
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "fallback_hash.safetensors")
|
||
with open(file_path, "wb") as f:
|
||
f.write(payload)
|
||
|
||
self.assertIsNone(extract_safetensors_header_sha256(file_path))
|
||
self.assertEqual(calculate_file_sha256(file_path), expected_hash)
|
||
|
||
def test_extract_response_file_size_html_type_exclusion(self):
|
||
from requests.structures import CaseInsensitiveDict
|
||
|
||
# HTML Content-Type should be ignored for content-length
|
||
mock_response = MagicMock()
|
||
mock_response.headers = CaseInsensitiveDict({
|
||
"Content-Length": "12345",
|
||
"Content-Type": "text/html; charset=utf-8"
|
||
})
|
||
size = extract_response_file_size(mock_response)
|
||
self.assertIsNone(size)
|
||
|
||
# Non-HTML content-type should be parsed
|
||
mock_response_valid = MagicMock()
|
||
mock_response_valid.headers = CaseInsensitiveDict({
|
||
"Content-Length": "12345",
|
||
"Content-Type": "application/octet-stream"
|
||
})
|
||
size_valid = extract_response_file_size(mock_response_valid)
|
||
self.assertEqual(size_valid, 12345)
|
||
|
||
def test_read_json_safe_edge_cases(self):
|
||
# 1. Path is None
|
||
self.assertIsNone(read_json_safe(None))
|
||
# 2. Path is empty
|
||
self.assertEqual(read_json_safe("", "default"), "default")
|
||
# 3. Path is a directory
|
||
import tempfile
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
self.assertEqual(read_json_safe(tmpdir, "default"), "default")
|
||
|
||
def test_get_version_sort_key_non_dict_input(self):
|
||
# If version is not a dict, it must fall back to ("", 0) gracefully
|
||
self.assertEqual(get_version_sort_key(None), ("", 0))
|
||
self.assertEqual(get_version_sort_key("not a dict"), ("", 0))
|
||
self.assertEqual(get_version_sort_key([]), ("", 0))
|
||
|
||
def test_looks_like_model_file_expected_filename_fallback(self):
|
||
# expected_filename does not match basename, but it has a valid extension
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/model.safetensors", "other_name.safetensors"))
|
||
# expected_filename does not match basename, and does NOT have a valid extension
|
||
self.assertFalse(looks_like_model_file("https://example.com/files/config.json", "other_name.safetensors"))
|
||
|
||
def test_normalize_model_image_edge_cases(self):
|
||
# 1. Non-dict input returns empty dict
|
||
self.assertEqual(normalize_model_image(None), {})
|
||
self.assertEqual(normalize_model_image([]), {})
|
||
|
||
# 2. default_civitai_url fallback works
|
||
raw_image = {
|
||
"url": "https://example.com/some_image.png",
|
||
}
|
||
normalized = normalize_model_image(raw_image, default_civitai_url="https://civitai.com/posts/fallback")
|
||
self.assertEqual(normalized["civitaiUrl"], "https://civitai.com/posts/fallback")
|
||
|
||
def test_metadata_sidecar_strict_credential_scrubbing(self):
|
||
"""
|
||
Covers Security Risk: API keys, cookies, or authorization tokens leaking into local JSON metadata sidecar files.
|
||
Verifies that build_model_resolver_metadata scrubs sensitive query parameters and headers from payloads.
|
||
"""
|
||
import json
|
||
from core.download.metadata import build_model_resolver_metadata
|
||
dest_path = "dummy_model.safetensors"
|
||
metadata = {
|
||
"download_url": "https://civitai.com/api/download/models/123?token=sensitivetoken123&session=cookie_val",
|
||
"source_url": "https://civitai.com/models/123?apiKey=sensitiveapikey123",
|
||
"headers": {"Authorization": "Bearer some-bearer-token"},
|
||
"hf_token": "hf_abc123secret",
|
||
}
|
||
|
||
payload = build_model_resolver_metadata(
|
||
dest_path=dest_path,
|
||
metadata=metadata,
|
||
category="loras",
|
||
source_url="https://civitai.com/api/download/models/123?token=sensitivetoken123"
|
||
)
|
||
|
||
serialized = json.dumps(payload)
|
||
self.assertNotIn("sensitivetoken123", serialized)
|
||
self.assertNotIn("cookie_val", serialized)
|
||
self.assertNotIn("sensitiveapikey123", serialized)
|
||
self.assertNotIn("hf_abc123secret", serialized)
|
||
self.assertNotIn("Authorization", payload)
|
||
self.assertNotIn("headers", payload)
|
||
|
||
def test_write_json_atomic_cleanup_on_failure(self):
|
||
"""
|
||
Covers File System Integrity Risk: Disk write or rename failure leaving corrupted or incomplete files.
|
||
Verifies that write_json_atomic cleans up temporary files and leaves the original file intact on failure.
|
||
"""
|
||
import tempfile
|
||
import os
|
||
from unittest.mock import patch
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "target.json")
|
||
with open(file_path, "w", encoding="utf-8") as f:
|
||
f.write('{"intact": true}')
|
||
|
||
with patch("os.replace", side_effect=OSError("Permission denied")):
|
||
with self.assertRaises(OSError):
|
||
write_json_atomic(file_path, {"changed": True})
|
||
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
self.assertEqual(f.read(), '{"intact": true}')
|
||
|
||
files = os.listdir(tmpdir)
|
||
self.assertEqual(files, ["target.json"])
|
||
|
||
@patch("requests.head")
|
||
@patch("requests.get")
|
||
def test_fetch_remote_file_size_head_fails_get_range_succeeds(self, mock_get, mock_head):
|
||
"""
|
||
Covers Network Compatibility Risk: HEAD request blocked or unsupported by remote servers.
|
||
Verifies fallback to partial GET request using bytes=0-0 range and extracting correct Content-Range size.
|
||
Asserts that Range headers and stream parameters are correct to prevent full model downloads.
|
||
"""
|
||
from requests.structures import CaseInsensitiveDict
|
||
|
||
mock_head_res = MagicMock()
|
||
mock_head_res.status_code = 405
|
||
mock_head.return_value = mock_head_res
|
||
|
||
mock_get_res = MagicMock()
|
||
mock_get_res.status_code = 206
|
||
mock_get_res.headers = CaseInsensitiveDict({
|
||
"Content-Range": "bytes 0-0/987654321",
|
||
"Content-Length": "1",
|
||
"Content-Type": "application/octet-stream"
|
||
})
|
||
mock_get.return_value = mock_get_res
|
||
|
||
size = fetch_remote_file_size("https://example.com/range_model.safetensors", headers={"Custom-Header": "Value"})
|
||
self.assertEqual(size, 987654321)
|
||
mock_head.assert_called_once()
|
||
|
||
# Assert specific range check request parameters to prevent full file download
|
||
mock_get.assert_called_once_with(
|
||
"https://example.com/range_model.safetensors",
|
||
headers={"Custom-Header": "Value", "Accept-Encoding": "identity", "Range": "bytes=0-0"},
|
||
allow_redirects=False,
|
||
stream=True,
|
||
timeout=15
|
||
)
|
||
|
||
def test_calculate_file_sha256_handles_large_file_chunking(self):
|
||
"""
|
||
Covers Process Stability Risk: OOM error when hashing large gigabyte-sized files.
|
||
Verifies calculate_file_sha256 reads the file incrementally in exact chunk_size blocks,
|
||
and that a mutation reading the entire file in one go is caught.
|
||
"""
|
||
from unittest.mock import mock_open, patch
|
||
import hashlib
|
||
|
||
content = b"large_model_data_block_chunk_test"
|
||
expected_hash = hashlib.sha256(content).hexdigest()
|
||
|
||
# Setup a mock file where we can control read outputs and count calls
|
||
# We simulate reading 4 bytes at a time
|
||
chunks = [content[i:i+4] for i in range(0, len(content), 4)] + [b""]
|
||
mock_file = MagicMock()
|
||
mock_file.read.side_effect = chunks
|
||
mock_file.__enter__.return_value = mock_file
|
||
|
||
m_open = mock_open()
|
||
m_open.return_value = mock_file
|
||
|
||
with patch("builtins.open", m_open), patch("os.path.exists", return_value=True):
|
||
calculated_hash = calculate_file_sha256("dummy_path.safetensors", chunk_size=4)
|
||
|
||
# Assert correct hash was calculated
|
||
self.assertEqual(calculated_hash, expected_hash)
|
||
|
||
# Assert it was read chunk-by-chunk rather than all at once
|
||
self.assertEqual(mock_file.read.call_count, len(chunks))
|
||
for call in mock_file.read.call_args_list[:-1]:
|
||
self.assertEqual(call[0][0], 4)
|
||
|
||
def test_read_json_safe_handles_raw_corrupt_payload(self):
|
||
"""
|
||
Covers Resilience Risk: Malformed or truncated settings JSON causing process startup failure.
|
||
Verifies read_json_safe catches json formatting/syntax errors and successfully returns the fallback.
|
||
"""
|
||
import tempfile
|
||
import os
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "corrupt.json")
|
||
with open(file_path, "w", encoding="utf-8") as f:
|
||
f.write("corrupted JSON content {[[}")
|
||
|
||
result = read_json_safe(file_path, default={"fallback": "value"})
|
||
self.assertEqual(result, {"fallback": "value"})
|
||
|
||
def test_fetch_remote_file_size_ignores_html_content_type(self):
|
||
"""
|
||
Covers Network Response Boundary: Cloudflare or local proxies returning HTML pages with length instead of binaries.
|
||
Verifies extract_response_file_size excludes HTML page content-lengths from being parsed as valid size metadata.
|
||
"""
|
||
from requests.structures import CaseInsensitiveDict
|
||
mock_response = MagicMock()
|
||
mock_response.headers = CaseInsensitiveDict({
|
||
"Content-Length": "8888",
|
||
"Content-Type": "text/html; charset=utf-8"
|
||
})
|
||
|
||
size = extract_response_file_size(mock_response)
|
||
self.assertIsNone(size)
|
||
|
||
def test_get_version_sort_key_missing_or_invalid_values(self):
|
||
"""
|
||
Covers Functional Robustness: Missing or invalid timestamps/IDs in model versions.
|
||
Verifies stable key tuples are generated without raising TypeError or ValueError.
|
||
"""
|
||
version_missing_id = {"published_at": "2024-01-01T00:00:00Z"}
|
||
version_invalid_id = {"published_at": "2024-01-01T00:00:00Z", "id": "not_an_int"}
|
||
version_none = None
|
||
|
||
key_missing = get_version_sort_key(version_missing_id)
|
||
key_invalid = get_version_sort_key(version_invalid_id)
|
||
key_none = get_version_sort_key(version_none)
|
||
|
||
self.assertEqual(key_missing, ("2024-01-01T00:00:00Z", 0))
|
||
self.assertEqual(key_invalid, ("2024-01-01T00:00:00Z", 0))
|
||
self.assertEqual(key_none, ("", 0))
|
||
|
||
def test_looks_like_model_file_strict_matches(self):
|
||
"""
|
||
Covers Input Boundary Validation: Ensuring downloader model URL verification functions exclude false positives.
|
||
Verifies HuggingFace Spaces page rejection and valid download link acceptance.
|
||
"""
|
||
self.assertFalse(looks_like_model_file("https://huggingface.co/spaces/user/space-name"))
|
||
self.assertTrue(looks_like_model_file("https://civitai.com/api/download/models/12345"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/some_model.safetensors"))
|
||
self.assertFalse(looks_like_model_file("https://example.com/some_image.jpg"))
|
||
|
||
@patch("requests.get")
|
||
def test_check_credential_http_rate_limit_and_timeout(self, mock_get):
|
||
"""
|
||
Covers Integration Boundary Risk: Remote credential checking fails gracefully under rate limits or network timeout.
|
||
Verifies Timeout exception handling and HTTP 429 rate limit responses map correctly.
|
||
"""
|
||
import requests
|
||
mock_get.side_effect = requests.exceptions.Timeout("Request timed out")
|
||
|
||
result_timeout = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={"Authorization": "Bearer token"}
|
||
)
|
||
|
||
self.assertFalse(result_timeout["success"])
|
||
self.assertEqual(result_timeout["status"], "timeout")
|
||
|
||
mock_get.side_effect = None
|
||
mock_res_429 = MagicMock()
|
||
mock_res_429.status_code = 429
|
||
mock_get.return_value = mock_res_429
|
||
|
||
result_429 = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={"Authorization": "Bearer token"}
|
||
)
|
||
|
||
self.assertFalse(result_429["success"])
|
||
self.assertEqual(result_429["status_code"], 429)
|
||
self.assertEqual(result_429["status"], "error")
|
||
|
||
def test_write_json_atomic_serialization_failure_cleanup(self):
|
||
"""
|
||
Covers File System Integrity Risk: Attempting to serialize un-serializable payloads during atomic writes.
|
||
Verifies that any temp files generated before/during serialization failure are cleaned up.
|
||
"""
|
||
import tempfile
|
||
import os
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "target.json")
|
||
with open(file_path, "w", encoding="utf-8") as f:
|
||
f.write('{"intact": true}')
|
||
|
||
with self.assertRaises(TypeError):
|
||
write_json_atomic(file_path, {"unserializable": {1, 2, 3}})
|
||
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
self.assertEqual(f.read(), '{"intact": true}')
|
||
|
||
files = os.listdir(tmpdir)
|
||
self.assertEqual(files, ["target.json"])
|
||
|
||
def test_looks_like_model_file_encoded_and_fragments(self):
|
||
"""
|
||
Covers Input Boundary Validation: Validates looks_like_model_file parses percent-encoded URLs,
|
||
query parameters, and URL fragment/hash symbols correctly.
|
||
"""
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/model%2Esafetensors"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/model.ckpt#weights"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/files/my%20model.safetensors?token=abc"))
|
||
self.assertTrue(looks_like_model_file("https://example.com/my%20model.safetensors", "my model.safetensors"))
|
||
|
||
@patch("requests.get")
|
||
def test_check_credential_http_custom_429_handler_raises(self, mock_get):
|
||
"""
|
||
Covers Error Recovery Risk: Simulates custom rate limit handlers raising unhandled exceptions.
|
||
Verifies that check_credential_http catches these exceptions and maps them gracefully to failure responses.
|
||
"""
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 429
|
||
mock_get.return_value = mock_response
|
||
|
||
def bad_handler(resp):
|
||
raise ValueError("custom handler exploded")
|
||
|
||
result = check_credential_http(
|
||
url="https://civitai.com/api/v1/me",
|
||
headers={},
|
||
custom_429_handler=bad_handler
|
||
)
|
||
|
||
self.assertFalse(result["success"])
|
||
self.assertFalse(result["valid"])
|
||
self.assertEqual(result["status"], "error")
|
||
self.assertIn("custom handler exploded", result["message"])
|
||
|
||
def test_normalize_sha256(self):
|
||
h = "a" * 64
|
||
self.assertEqual(normalize_sha256(h), h)
|
||
self.assertEqual(normalize_sha256("sha256:" + h), h)
|
||
self.assertEqual(normalize_sha256("sha256=" + h), h)
|
||
self.assertEqual(normalize_sha256("SHA256:" + h.upper()), h)
|
||
self.assertEqual(normalize_sha256("0x" + h), h)
|
||
self.assertEqual(normalize_sha256("sha256:0x" + h), h)
|
||
self.assertEqual(normalize_sha256("SHA256:0X" + h.upper()), h)
|
||
self.assertEqual(normalize_sha256("invalid_hash"), "")
|
||
self.assertEqual(normalize_sha256(None), "")
|
||
self.assertEqual(normalize_sha256(123), "")
|
||
|
||
def test_unique_ordered_strings(self):
|
||
self.assertEqual(unique_ordered_strings(["a", "b", "a", "c", "b"]), ["a", "b", "c"])
|
||
self.assertEqual(unique_ordered_strings([" a ", "b", "a"]), ["a", "b"])
|
||
self.assertEqual(unique_ordered_strings([]), [])
|
||
|
||
def test_extract_sha256_from_metadata(self):
|
||
h = "a" * 64
|
||
self.assertEqual(extract_sha256_from_metadata({"sha256": h}), h)
|
||
self.assertEqual(extract_sha256_from_metadata({"hash": h}), h)
|
||
self.assertEqual(extract_sha256_from_metadata({"hashes": {"SHA256": h}}), h)
|
||
self.assertEqual(extract_sha256_from_metadata({"hashes": {"sha256": h}}), h)
|
||
self.assertEqual(extract_sha256_from_metadata({"file_info": {"sha256": h}}), h)
|
||
self.assertEqual(extract_sha256_from_metadata({"file_info": {"hashes": {"sha256": h}}}), h)
|
||
self.assertEqual(extract_sha256_from_metadata({}), "")
|
||
|
||
def test_extract_trained_words(self):
|
||
# --- basic flat inputs (existing cases) ---
|
||
self.assertEqual(extract_trained_words("word1", ["word2"]), ["word1", "word2"])
|
||
self.assertEqual(extract_trained_words({"word": "word1"}, [{"name": "word2"}]), ["word1", "word2"])
|
||
self.assertEqual(extract_trained_words("word1", "word1"), ["word1"])
|
||
|
||
def test_extract_trained_words_version_dict_trainedWords(self):
|
||
# Mirrors the real civitai.py call: extract_trained_words(version)
|
||
# where version is a dict with "trainedWords" key (list of strings)
|
||
version = {
|
||
"id": 123,
|
||
"name": "v1.0",
|
||
"trainedWords": ["trigger1", "trigger2"],
|
||
"baseModel": "SD 1.5",
|
||
}
|
||
self.assertEqual(extract_trained_words(version), ["trigger1", "trigger2"])
|
||
|
||
def test_extract_trained_words_version_dict_trainedWords_string(self):
|
||
# trainedWords may be a single string in some API responses
|
||
version = {"trainedWords": "solo_trigger"}
|
||
self.assertEqual(extract_trained_words(version), ["solo_trigger"])
|
||
|
||
def test_extract_trained_words_version_dict_trigger_key(self):
|
||
# Mirrors the civarchive.py call pattern: version.get("trigger")
|
||
version = {"trigger": "civarchive_trigger"}
|
||
self.assertEqual(extract_trained_words(version), ["civarchive_trigger"])
|
||
|
||
def test_extract_trained_words_version_dict_both_trainedWords_and_trigger(self):
|
||
# Both keys present – both should be collected, deduped
|
||
version = {"trainedWords": ["word_a", "word_b"], "trigger": "word_b"}
|
||
result = extract_trained_words(version)
|
||
self.assertIn("word_a", result)
|
||
self.assertIn("word_b", result)
|
||
self.assertEqual(result.count("word_b"), 1) # deduplication
|
||
|
||
def test_extract_trained_words_version_dict_model_tags(self):
|
||
# model.tags path
|
||
version = {"model": {"tags": ["tag1", "tag2"]}}
|
||
self.assertEqual(extract_trained_words(version), ["tag1", "tag2"])
|
||
|
||
def test_extract_trained_words_empty_and_none_inputs(self):
|
||
self.assertEqual(extract_trained_words(None), [])
|
||
self.assertEqual(extract_trained_words(""), [])
|
||
self.assertEqual(extract_trained_words({}), [])
|
||
self.assertEqual(extract_trained_words({"trainedWords": []}), [])
|
||
self.assertEqual(extract_trained_words({"trainedWords": None}), [])
|
||
|
||
def test_extract_trained_words_deduplication_case_insensitive(self):
|
||
# Same word in different cases should appear only once
|
||
result = extract_trained_words(["Trigger1", "trigger1", "TRIGGER1"])
|
||
self.assertEqual(len(result), 1)
|
||
self.assertEqual(result[0], "Trigger1") # first occurrence wins
|
||
|
||
def test_extract_trained_words_multiple_version_dicts(self):
|
||
# Called with multiple separate version objects (multi-arg form)
|
||
v1 = {"trainedWords": ["alpha"]}
|
||
v2 = {"trainedWords": ["beta"]}
|
||
self.assertEqual(extract_trained_words(v1, v2), ["alpha", "beta"])
|
||
|
||
@patch("core.type_utils.fetch_remote_file_size")
|
||
def test_fetch_remote_file_size_cached(self, mock_fetch):
|
||
clear_remote_size_cache()
|
||
mock_fetch.return_value = 12345
|
||
url = "https://example.com/file.bin"
|
||
|
||
size1 = fetch_remote_file_size_cached(url)
|
||
size2 = fetch_remote_file_size_cached(url)
|
||
|
||
self.assertEqual(size1, 12345)
|
||
self.assertEqual(size2, 12345)
|
||
mock_fetch.assert_called_once()
|
||
|
||
# Test clearing cache
|
||
clear_remote_size_cache()
|
||
mock_fetch.reset_mock()
|
||
mock_fetch.return_value = 54321
|
||
size3 = fetch_remote_file_size_cached(url)
|
||
self.assertEqual(size3, 54321)
|
||
mock_fetch.assert_called_once()
|
||
|
||
@patch("os.path.exists")
|
||
def test_find_metadata_sidecar_path(self, mock_exists):
|
||
# We want to test finding a sidecar file.
|
||
# Let's say model path is e:/models/sd15.safetensors
|
||
# Standard names: e:/models/sd15.metadata.json, e:/models/sd15.safetensors.metadata.json, etc.
|
||
mock_exists.side_effect = lambda p: p.endswith("sd15.metadata.json")
|
||
|
||
path = find_metadata_sidecar_path("e:/models/sd15.safetensors")
|
||
self.assertTrue(path.endswith("sd15.metadata.json"))
|
||
|
||
@patch("os.path.exists")
|
||
def test_find_metadata_sidecar_path_does_not_return_input_path(self, mock_exists):
|
||
mock_exists.side_effect = (
|
||
lambda p: p.replace("\\", "/").lower() == "e:/models/ae.metadata.json"
|
||
)
|
||
|
||
path = find_metadata_sidecar_path("e:/models/ae.metadata.json")
|
||
self.assertEqual(path, "")
|
||
|
||
def test_sidecar_candidates_prefer_exact_names_when_both_forms_exist(self):
|
||
import os
|
||
import tempfile
|
||
|
||
from core.path_utils import (
|
||
_metadata_sidecar_paths,
|
||
find_external_metadata_sidecar_path,
|
||
)
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
model_path = os.path.join(tmpdir, "model.safetensors")
|
||
legacy_path = os.path.join(tmpdir, "model.metadata.json")
|
||
exact_path = os.path.join(tmpdir, "model.safetensors.metadata.json")
|
||
for path in (model_path, legacy_path, exact_path):
|
||
with open(path, "w", encoding="utf-8") as handle:
|
||
handle.write("{}")
|
||
|
||
self.assertEqual(_metadata_sidecar_paths(model_path)[0], exact_path)
|
||
self.assertEqual(find_external_metadata_sidecar_path(model_path), exact_path)
|
||
|
||
def test_external_legacy_sidecar_keeps_collision_matching_rules(self):
|
||
import json
|
||
import os
|
||
import tempfile
|
||
|
||
from core.path_utils import find_external_metadata_sidecar_path
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
model_path = os.path.join(tmpdir, "model.safetensors")
|
||
collision_path = os.path.join(tmpdir, "model.ckpt")
|
||
legacy_path = os.path.join(tmpdir, "model.metadata.json")
|
||
for path in (model_path, collision_path):
|
||
with open(path, "wb") as handle:
|
||
handle.write(b"model")
|
||
|
||
with open(legacy_path, "w", encoding="utf-8") as handle:
|
||
json.dump({"filename": "model.ckpt"}, handle)
|
||
self.assertEqual(find_external_metadata_sidecar_path(model_path), "")
|
||
|
||
with open(legacy_path, "w", encoding="utf-8") as handle:
|
||
json.dump({"filename": "model.safetensors"}, handle)
|
||
self.assertEqual(find_external_metadata_sidecar_path(model_path), legacy_path)
|
||
|
||
def test_merges_external_user_fields_with_model_resolver_metadata(self):
|
||
import json
|
||
import os
|
||
import tempfile
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
model_path = os.path.join(tmpdir, "layered.safetensors")
|
||
with open(model_path, "wb") as handle:
|
||
handle.write(b"model")
|
||
|
||
external_path = os.path.join(tmpdir, "layered.metadata.json")
|
||
external_payload = {
|
||
"notes": "manual note",
|
||
"favorite": True,
|
||
"description": "manual description",
|
||
"tags": ["external", "shared"],
|
||
"trained_words": ["manual trigger"],
|
||
"sha256": "a" * 64,
|
||
}
|
||
with open(external_path, "w", encoding="utf-8") as handle:
|
||
json.dump(external_payload, handle)
|
||
|
||
resolver_path = get_model_resolver_sidecar_path(model_path)
|
||
resolver_payload = {
|
||
"managed_by": "comfyui-model-resolver",
|
||
"notes": "must not replace external note",
|
||
"description": "remote description",
|
||
"tags": ["resolver", "shared"],
|
||
"trained_words": ["remote trigger"],
|
||
"sha256": "b" * 64,
|
||
}
|
||
with open(resolver_path, "w", encoding="utf-8") as handle:
|
||
json.dump(resolver_payload, handle)
|
||
|
||
merged = read_merged_model_metadata(model_path, {})
|
||
selected_path = find_metadata_sidecar_path(model_path)
|
||
with open(external_path, "r", encoding="utf-8") as handle:
|
||
external_after = json.load(handle)
|
||
|
||
self.assertEqual(resolver_path, selected_path)
|
||
self.assertEqual(external_payload, external_after)
|
||
self.assertEqual("b" * 64, merged["sha256"])
|
||
self.assertEqual("manual note", merged["notes"])
|
||
self.assertTrue(merged["favorite"])
|
||
self.assertEqual("manual description", merged["description"])
|
||
self.assertEqual(["external", "shared", "resolver"], merged["tags"])
|
||
self.assertEqual(
|
||
["manual trigger", "remote trigger"],
|
||
merged["trained_words"],
|
||
)
|
||
|
||
def test_normalize_category_to_model_type(self):
|
||
self.assertEqual(normalize_category_to_model_type("checkpoints"), "checkpoint")
|
||
self.assertEqual(normalize_category_to_model_type("loras"), "lora")
|
||
self.assertEqual(normalize_category_to_model_type("vae"), "vae")
|
||
self.assertEqual(normalize_category_to_model_type("diffusion_models"), "diffusion_model")
|
||
self.assertEqual(normalize_category_to_model_type("embeddings"), "embedding")
|
||
self.assertEqual(normalize_category_to_model_type(""), "")
|
||
|
||
def test_check_credential_preconditions(self):
|
||
# Empty inputs
|
||
self.assertEqual(
|
||
check_credential_preconditions("", "Token"),
|
||
{
|
||
"success": False,
|
||
"valid": False,
|
||
"status": "missing",
|
||
"message": "Paste a Token first.",
|
||
}
|
||
)
|
||
# Non-empty inputs
|
||
self.assertIsNone(check_credential_preconditions("my-token", "Token"))
|
||
|
||
def test_calculate_file_sha256_with_progress(self):
|
||
import tempfile
|
||
import os
|
||
from core.path_utils import calculate_file_sha256
|
||
|
||
content = b"some test binary data for progress"
|
||
progress_calls = []
|
||
|
||
def on_progress(bytes_read, total_bytes):
|
||
progress_calls.append((bytes_read, total_bytes))
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "progress_test.bin")
|
||
with open(file_path, "wb") as f:
|
||
f.write(content)
|
||
|
||
calculated_hash = calculate_file_sha256(
|
||
file_path, chunk_size=4, on_progress=on_progress
|
||
)
|
||
self.assertIsNotNone(calculated_hash)
|
||
# Verify progress was called
|
||
self.assertTrue(len(progress_calls) > 0)
|
||
self.assertEqual(progress_calls[-1], (len(content), len(content)))
|
||
|
||
def test_calculate_file_sha256_cancellation(self):
|
||
import tempfile
|
||
import os
|
||
from core.path_utils import calculate_file_sha256, HashCalculationCancelled
|
||
|
||
content = b"some test binary data for cancellation"
|
||
|
||
def is_cancelled():
|
||
return True
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
file_path = os.path.join(tmpdir, "cancel_test.bin")
|
||
with open(file_path, "wb") as f:
|
||
f.write(content)
|
||
|
||
with self.assertRaises(HashCalculationCancelled):
|
||
calculate_file_sha256(
|
||
file_path, chunk_size=4, is_cancelled=is_cancelled
|
||
)
|
||
|
||
def test_save_catalog_with_backup(self):
|
||
import tempfile
|
||
import os
|
||
from core.path_utils import save_catalog_with_backup, read_json_safe
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
data_file = os.path.join(tmpdir, "catalog.json")
|
||
meta_file = os.path.join(tmpdir, "catalog_meta.json")
|
||
|
||
# Initial write
|
||
initial_data = {"items": [1, 2, 3]}
|
||
initial_meta = {"updated_at": "now"}
|
||
|
||
save_catalog_with_backup(data_file, initial_data, meta_file, initial_meta)
|
||
|
||
self.assertTrue(os.path.exists(data_file))
|
||
self.assertTrue(os.path.exists(meta_file))
|
||
self.assertEqual(read_json_safe(data_file), initial_data)
|
||
self.assertEqual(read_json_safe(meta_file), initial_meta)
|
||
|
||
# Write with backup check
|
||
new_data = {"items": [4, 5, 6]}
|
||
new_meta = {"updated_at": "later"}
|
||
|
||
save_catalog_with_backup(data_file, new_data, meta_file, new_meta)
|
||
|
||
# Verify backup exists
|
||
backup_file = data_file + ".bak"
|
||
self.assertTrue(os.path.exists(backup_file))
|
||
self.assertEqual(read_json_safe(backup_file), initial_data)
|
||
|
||
# Verify new data is written
|
||
self.assertEqual(read_json_safe(data_file), new_data)
|
||
self.assertEqual(read_json_safe(meta_file), new_meta)
|
||
|
||
@patch("requests.get")
|
||
def test_request_source_json_success(self, mock_get):
|
||
from core.network_utils import request_source_json
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.json.return_value = {"success": True}
|
||
mock_get.return_value = mock_response
|
||
|
||
res = request_source_json("https://example.com/api")
|
||
self.assertEqual(res, {"success": True})
|
||
|
||
@patch("requests.get")
|
||
def test_request_source_json_accepts_array_payload(self, mock_get):
|
||
from core.network_utils import request_source_json
|
||
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.json.return_value = [{"id": "author/model"}]
|
||
mock_get.return_value = mock_response
|
||
|
||
res = request_source_json("https://example.com/api")
|
||
|
||
self.assertEqual(res, [{"id": "author/model"}])
|
||
|
||
@patch("requests.get")
|
||
@patch("time.sleep")
|
||
def test_request_source_json_429_retry(self, mock_sleep, mock_get):
|
||
from core.network_utils import request_source_json
|
||
mock_response_429 = MagicMock()
|
||
mock_response_429.status_code = 429
|
||
mock_response_429.headers = {"Retry-After": "1.5"}
|
||
|
||
mock_response_200 = MagicMock()
|
||
mock_response_200.status_code = 200
|
||
mock_response_200.json.return_value = {"ok": True}
|
||
|
||
mock_get.side_effect = [mock_response_429, mock_response_200]
|
||
|
||
res = request_source_json("https://example.com/api", max_attempts=2)
|
||
self.assertEqual(res, {"ok": True})
|
||
mock_sleep.assert_called_once_with(1.5)
|
||
|
||
def test_split_path_segments(self):
|
||
from core.path_utils import split_path_segments
|
||
self.assertEqual(split_path_segments("a/b/c"), ["a", "b", "c"])
|
||
self.assertEqual(split_path_segments("a\\b\\c"), ["a", "b", "c"])
|
||
self.assertEqual(split_path_segments("a/./b/../c"), ["a", "b", "c"])
|
||
self.assertEqual(split_path_segments("a/./b/../c", filter_dots=False), ["a", ".", "b", "..", "c"])
|
||
self.assertEqual(split_path_segments(""), [])
|
||
self.assertEqual(split_path_segments(None), [])
|
||
|
||
def test_parse_provider_model_url(self):
|
||
from core.type_utils import parse_provider_model_url
|
||
allowed = ["civitai.com", "civitai.red"]
|
||
|
||
# Valid url download format
|
||
res = parse_provider_model_url("https://civitai.com/api/download/models/12345", allowed)
|
||
self.assertEqual(res.to_dict(), {"version_id": 12345})
|
||
|
||
# Valid model page format
|
||
res = parse_provider_model_url("https://civitai.red/models/9876?modelVersionId=54321", allowed)
|
||
self.assertEqual(
|
||
res.to_dict(),
|
||
{"model_id": 9876, "version_id": 54321},
|
||
)
|
||
|
||
# Non-matching hostname
|
||
res = parse_provider_model_url("https://example.com/api/download/models/12345", allowed)
|
||
self.assertIsNone(res)
|
||
|
||
# Mirror sha256 path pattern
|
||
res = parse_provider_model_url("https://civarchive.com/sha256/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2", ["civarchive.com"])
|
||
self.assertIsNotNone(res)
|
||
self.assertEqual(len(res.sha256), 64)
|
||
self.assertEqual(res.to_dict(), {"sha256": res.sha256})
|
||
|
||
def test_resolve_model_category(self):
|
||
from core.type_utils import resolve_model_category
|
||
self.assertEqual(resolve_model_category("checkpoints", target_format="folder"), "checkpoints")
|
||
self.assertEqual(resolve_model_category("lora", target_format="folder"), "loras")
|
||
self.assertEqual(resolve_model_category("select_safetensors", target_format="folder"), "diffusion_models")
|
||
self.assertEqual(resolve_model_category("checkpoint", target_format="civitai"), "Checkpoint")
|
||
self.assertEqual(resolve_model_category("diffusion_models", target_format="civitai"), "Checkpoint")
|
||
self.assertEqual(resolve_model_category("loras", target_format="civitai"), "LORA")
|
||
self.assertEqual(resolve_model_category("locon", target_format="civarchive"), "LoCon")
|
||
self.assertEqual(resolve_model_category("loras", target_format="urn"), "loras")
|
||
|
||
@patch("core.network_utils.request_source_json")
|
||
def test_execute_provider_json_request(self, mock_request_json):
|
||
from core.network_utils import execute_provider_json_request
|
||
mock_request_json.return_value = {"items": []}
|
||
res = execute_provider_json_request(
|
||
"CivitAI API",
|
||
"https://civitai.com/api/v1/models",
|
||
params={"limit": 5},
|
||
session_token="secret_token",
|
||
)
|
||
self.assertEqual(res, {"items": []})
|
||
mock_request_json.assert_called_once()
|
||
call_kwargs = mock_request_json.call_args[1]
|
||
self.assertIn("Cookie", call_kwargs["headers"])
|
||
|
||
@patch("core.network_utils.request_source_json")
|
||
def test_execute_provider_json_request_accepts_array_payload(
|
||
self, mock_request_json
|
||
):
|
||
from core.network_utils import execute_provider_json_request
|
||
|
||
mock_request_json.return_value = [{"id": "author/model"}]
|
||
|
||
result = execute_provider_json_request(
|
||
"HuggingFace API",
|
||
"https://huggingface.co/api/models",
|
||
)
|
||
|
||
self.assertEqual(result, [{"id": "author/model"}])
|