import base64
import unittest
from diffusion_cli.config import (
GenerationDefaults,
ModelPathConfig,
UserConfig,
)
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.generation_service import GeneratedImage
from diffusion_cli.server import ServerConfig, createApp
class FakeGenerationService:
def __init__(self):
self.user_config = UserConfig(
models=ModelPathConfig(),
generation=GenerationDefaults(
width=512,
steps=12,
output_extension="jpeg",
output_quality=85,
),
)
self.requests = []
self.error = None
self.image = GeneratedImage(b"\x89PNG\r\n\x1a\nimage", "png", 321)
def generateImages(self, request):
if self.error:
raise self.error
self.requests.append(request)
return [self.image]
class UiTest(unittest.TestCase):
def makeClient(self):
service = FakeGenerationService()
app = createApp(
ServerConfig("sillytavern-sdcpp", "127.0.0.1", 7860, "cpu-cache"),
service,
)
return app.test_client(), service
def testUiHealthReturnsServerProfileInformation(self):
client, _service = self.makeClient()
response = client.get("/ui/api/health")
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json["status"], "ok")
self.assertEqual(response.json["api_profile"], "sillytavern-sdcpp")
self.assertEqual(response.json["model_residency"], "cpu-cache")
def testUiDefaultsComeFromGenerationConfig(self):
client, _service = self.makeClient()
response = client.get("/ui/api/defaults")
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json["negative_prompt"], "")
self.assertEqual(response.json["width"], 512)
self.assertEqual(response.json["height"], 1248)
self.assertEqual(response.json["batch_size"], 1)
self.assertEqual(response.json["steps"], 12)
self.assertEqual(response.json["cfg"], 1.0)
self.assertEqual(response.json["output_extension"], "jpg")
self.assertEqual(response.json["output_quality"], 85)
self.assertEqual(
response.json["output_formats"],
[
{"value": "png", "label": "PNG"},
{"value": "jpg", "label": "JPEG"},
{"value": "webp", "label": "WebP"},
{"value": "avif", "label": "AVIF"},
],
)
self.assertEqual(response.json["seed"], -1)
presets = response.json["resolution_presets"]
self.assertEqual(len(presets), 33)
self.assertEqual(
presets[0],
{
"width": 1024,
"height": 1024,
"resolution": "1024",
"aspect_ratio": "1:1",
"label": "1024x1024 (1:1)",
},
)
self.assertEqual(presets[-1]["label"], "864x2016 (9:21)")
def testUiStaticAssetsAreServedByExplicitRoutes(self):
client, _service = self.makeClient()
index = client.get("/ui/")
app_js = client.get("/ui/app.js")
style = client.get("/ui/style.css")
preact = client.get("/ui/vendor/preact.module.js")
self.assertEqual(index.status_code, 200)
self.assertIn(b'type="module"', index.data)
self.assertEqual(app_js.status_code, 200)
self.assertIn(b'from "/ui/vendor/preact.module.js"', app_js.data)
self.assertEqual(style.status_code, 200)
self.assertEqual(preact.status_code, 200)
self.assertIn(b"export", preact.data)
index.close()
app_js.close()
style.close()
preact.close()
def testUiRootRedirectsToSlashRoot(self):
client, _service = self.makeClient()
response = client.get("/ui")
self.assertEqual(response.status_code, 302)
self.assertEqual(response.headers["Location"], "/ui/")
def testUiTxt2ImgReturnsImageObjects(self):
client, service = self.makeClient()
response = client.post(
"/ui/api/txt2img",
json={
"prompt": " a mug ",
"negative_prompt": "",
"width": 512,
"height": 768,
"steps": 12,
"cfg": 1.5,
"seed": -1,
"batch_size": 2,
"output_extension": "jpg",
"output_quality": 85,
},
)
self.assertEqual(response.status_code, 200)
request = service.requests[0]
self.assertEqual(request.prompt, "a mug")
self.assertIsNone(request.negative_prompt)
self.assertEqual(request.width, 512)
self.assertEqual(request.height, 768)
self.assertEqual(request.steps, 12)
self.assertEqual(request.cfg, 1.5)
self.assertIsNone(request.seed)
self.assertEqual(request.batch_size, 2)
self.assertEqual(request.output_extension, "jpg")
self.assertEqual(request.output_quality, 85)
image = response.json["images"][0]
self.assertEqual(image["mime_type"], "image/png")
self.assertEqual(
base64.b64decode(image["data"]),
b"\x89PNG\r\n\x1a\nimage",
)
self.assertEqual(image["seed"], 321)
def testUiTxt2ImgUsesGeneratedImageMimeType(self):
client, service = self.makeClient()
service.image = GeneratedImage(b"jpeg", "jpg", 321)
response = client.post("/ui/api/txt2img", json={"prompt": "a mug"})
self.assertEqual(response.status_code, 200)
image = response.json["images"][0]
self.assertEqual(image["mime_type"], "image/jpeg")
self.assertEqual(base64.b64decode(image["data"]), b"jpeg")
def testUiTxt2ImgRejectsEmptyPrompt(self):
client, _service = self.makeClient()
response = client.post("/ui/api/txt2img", json={"prompt": " "})
self.assertEqual(response.status_code, 400)
self.assertEqual(response.json["error"], "bad_request")
self.assertEqual(response.json["message"], "prompt required")
def testUiTxt2ImgRejectsInvalidCfg(self):
client, _service = self.makeClient()
response = client.post(
"/ui/api/txt2img",
json={"prompt": "a mug", "cfg": -1},
)
self.assertEqual(response.status_code, 400)
self.assertIn("CFG must be non-negative", response.json["message"])
def testUiTxt2ImgRejectsInvalidOutputFieldTypes(self):
client, _service = self.makeClient()
response = client.post(
"/ui/api/txt2img",
json={"prompt": "a mug", "output_extension": 12},
)
self.assertEqual(response.status_code, 400)
self.assertEqual(
response.json["message"],
"output_extension must be a string",
)
response = client.post(
"/ui/api/txt2img",
json={"prompt": "a mug", "output_quality": True},
)
self.assertEqual(response.status_code, 400)
self.assertEqual(
response.json["message"],
"output_quality must be an integer",
)
def testUiTxt2ImgRejectsInvalidOutputValues(self):
client, _service = self.makeClient()
response = client.post(
"/ui/api/txt2img",
json={"prompt": "a mug", "output_extension": "gif"},
)
self.assertEqual(response.status_code, 400)
self.assertIn("Output extension must be", response.json["message"])
response = client.post(
"/ui/api/txt2img",
json={"prompt": "a mug", "output_quality": 0},
)
self.assertEqual(response.status_code, 400)
self.assertIn("Output quality must be", response.json["message"])
def testUiTxt2ImgRejectsZeroWidth(self):
client, _service = self.makeClient()
response = client.post(
"/ui/api/txt2img",
json={"prompt": "a mug", "width": 0},
)
self.assertEqual(response.status_code, 400)
self.assertIn("Width must be", response.json["message"])
def testUiTxt2ImgReportsDiffusionCliErrors(self):
client, service = self.makeClient()
service.error = DiffusionCliError("missing model")
response = client.post("/ui/api/txt2img", json={"prompt": "a mug"})
self.assertEqual(response.status_code, 500)
self.assertEqual(response.json["error"], "generation_failed")
self.assertEqual(response.json["message"], "missing model")
if __name__ == "__main__":
unittest.main()