1 files changed, 29 insertions, 0 deletions
diff --git a/test_generate.py b/test_generate.py
new file mode 100644
index 0000000..c2425d5
--- /dev/null
+++ b/test_generate.py
@@ -0,0 +1,29 @@
+import generate
+
+
+def test_registry_integrity():
+    assert generate.DEFAULT_MODEL in generate.MODELS
+    for key, repo in generate.MODELS.items():
+        assert isinstance(key, str) and key
+        assert isinstance(repo, str) and "/" in repo
+
+
+def test_resolve_known_key():
+    assert generate.resolve_model("sdxl") == generate.MODELS["sdxl"]
+    assert generate.resolve_model("realvis") == generate.MODELS["realvis"]
+
+
+def test_resolve_unknown_key_raises_listing_valid():
+    try:
+        generate.resolve_model("nope")
+    except SystemExit as e:
+        assert "sdxl" in str(e) and "realvis" in str(e) and "pony" in str(e)
+    else:
+        assert False, "expected SystemExit for unknown key"
+
+
+if __name__ == "__main__":
+    test_registry_integrity()
+    test_resolve_known_key()
+    test_resolve_unknown_key_raises_listing_valid()
+    print("ok")