144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
import unittest
|
|
|
|
from json_inference import fill_template_from_keywords, infer_profile_from_keywords
|
|
|
|
|
|
class JsonInferenceTests(unittest.TestCase):
|
|
def test_profile_inference_basics(self) -> None:
|
|
keywords = [
|
|
"zapthorn",
|
|
"electric",
|
|
"wolf",
|
|
"thunder",
|
|
"claw",
|
|
"speed",
|
|
"storm",
|
|
"agile",
|
|
"forest",
|
|
"recover",
|
|
"energy",
|
|
]
|
|
|
|
profile = infer_profile_from_keywords(keywords)
|
|
|
|
self.assertEqual(profile["name"], "Zapthorn")
|
|
self.assertEqual(profile["type"], "electric")
|
|
self.assertIn("thunder", profile["attacks"])
|
|
self.assertIn("claw", profile["attacks"])
|
|
self.assertIn("recover", profile["abilities"])
|
|
self.assertEqual(profile["habitat"], "forest")
|
|
self.assertGreaterEqual(profile["speed"], 80)
|
|
|
|
def test_fill_key_only_template(self) -> None:
|
|
template = {
|
|
"name": "",
|
|
"type": "",
|
|
"secondary_type": None,
|
|
"attacks": [],
|
|
"abilities": [],
|
|
"habitat": "",
|
|
"personality": [],
|
|
"hp": None,
|
|
"attack": None,
|
|
"defense": None,
|
|
"speed": None,
|
|
"description": "",
|
|
"keywords": [],
|
|
}
|
|
|
|
keywords = [
|
|
"furret",
|
|
"normal",
|
|
"tail",
|
|
"smash",
|
|
"tunnel",
|
|
"agile",
|
|
"cheerful",
|
|
"explore",
|
|
"endurance",
|
|
]
|
|
|
|
result = fill_template_from_keywords(template, keywords)
|
|
|
|
self.assertEqual(result["name"], "Furret")
|
|
self.assertEqual(result["type"], "normal")
|
|
self.assertIn("smash", result["attacks"])
|
|
self.assertIn("explore", result["abilities"])
|
|
self.assertEqual(result["habitat"], "tunnel")
|
|
self.assertIn("cheerful", result["personality"])
|
|
self.assertIsInstance(result["description"], str)
|
|
self.assertGreater(len(result["description"]), 20)
|
|
|
|
def test_fill_tcg_style_template(self) -> None:
|
|
template = {
|
|
"category": "Pokemon",
|
|
"name": "",
|
|
"hp": "",
|
|
"types": [""],
|
|
"description": "",
|
|
"stage": "",
|
|
"attacks": [
|
|
{"cost": [""], "name": "", "effect": ""},
|
|
{"cost": [""], "name": "", "effect": "", "damage": 0},
|
|
],
|
|
"weaknesses": [{"type": "", "value": ""}],
|
|
"retreat": 0,
|
|
}
|
|
|
|
keywords = [
|
|
"zapthorn",
|
|
"electric",
|
|
"thunder",
|
|
"claw",
|
|
"speed",
|
|
"storm",
|
|
"energy",
|
|
]
|
|
|
|
result = fill_template_from_keywords(template, keywords)
|
|
|
|
self.assertEqual(result["name"], "Zapthorn")
|
|
self.assertEqual(result["types"], ["electric"])
|
|
self.assertEqual(result["stage"], "Basic")
|
|
self.assertTrue(result["hp"].isdigit())
|
|
self.assertEqual(result["weaknesses"][0]["type"], "ground")
|
|
self.assertEqual(result["weaknesses"][0]["value"], "x2")
|
|
self.assertEqual(result["attacks"][0]["name"], "Thunder")
|
|
self.assertEqual(result["attacks"][1]["name"], "Claw")
|
|
self.assertEqual(result["attacks"][0]["cost"], ["Electric"])
|
|
self.assertGreaterEqual(result["retreat"], 0)
|
|
|
|
def test_name_fallback_to_unknown_for_generic_tokens(self) -> None:
|
|
keywords = [
|
|
"black",
|
|
"fur",
|
|
"giant",
|
|
"electric",
|
|
"claw",
|
|
"speed",
|
|
"storm",
|
|
]
|
|
|
|
profile = infer_profile_from_keywords(keywords)
|
|
self.assertEqual(profile["name"], "Unknown")
|
|
|
|
def test_preserves_existing_values(self) -> None:
|
|
template = {
|
|
"name": "CustomName",
|
|
"type": "electric",
|
|
"attacks": [],
|
|
"description": "Already set",
|
|
}
|
|
keywords = ["furret", "normal", "attack"]
|
|
|
|
result = fill_template_from_keywords(template, keywords)
|
|
|
|
self.assertEqual(result["name"], "CustomName")
|
|
self.assertEqual(result["type"], "electric")
|
|
self.assertEqual(result["description"], "Already set")
|
|
self.assertIn("attack", result["attacks"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|