Files
AEmotionStudio-ComfyUI-Disc…/tests/test_message_builder.py
T
AEmotionStudioandClaude Opus 4.5 436b773bd8 chore: implement phase 6 - final cleanup
- Remove unused imports from nodes (time, torch, requests, itertools)
- Move uuid import to top-level in bot/bot.py
- Add 4 new test files with 86 tests for shared utilities
- Update CHANGELOG.md with version 2.0.0 refactoring summary
- Update REFACTOR_PRD.md to mark Phase 5 complete

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-20 10:20:51 -08:00

262 lines
9.1 KiB
Python

"""Tests for shared/discord/message_builder.py"""
import sys
import os
import unittest
from unittest.mock import MagicMock
# Mock dependencies before importing project modules
sys.modules["torch"] = MagicMock()
sys.modules["numpy"] = MagicMock()
sys.modules["cv2"] = MagicMock()
# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from shared.discord.message_builder import (
build_metadata_section,
build_prompt_section,
build_discord_message,
validate_message_content,
format_file_info,
format_file_size,
)
class TestBuildMetadataSection(unittest.TestCase):
"""Test build_metadata_section function."""
def test_empty_dict(self):
"""Test with empty info dict returns empty string."""
result = build_metadata_section({})
self.assertEqual(result, "")
def test_with_date(self):
"""Test metadata with date."""
result = build_metadata_section({"date": "2026-01-20"})
self.assertIn("**Date:** 2026-01-20", result)
self.assertIn("**Information:**", result)
def test_with_time(self):
"""Test metadata with time."""
result = build_metadata_section({"time": "14-30-00"})
self.assertIn("**Time:** 14-30-00", result)
def test_with_dimensions(self):
"""Test metadata with dimensions."""
result = build_metadata_section({"dimensions": "1920x1080"})
self.assertIn("**Dimensions:** 1920x1080", result)
def test_with_format(self):
"""Test metadata with file format."""
result = build_metadata_section({}, file_format="png")
self.assertIn("**Format:** PNG", result)
def test_with_frame_rate(self):
"""Test metadata with frame rate."""
result = build_metadata_section({}, frame_rate=30.0)
self.assertIn("**Frame Rate:** 30.0 fps", result)
def test_custom_section_title(self):
"""Test custom section title."""
result = build_metadata_section({"date": "2026-01-20"}, section_title="Video Info")
self.assertIn("**Video Info:**", result)
def test_exclude_options(self):
"""Test excluding certain metadata."""
info = {"date": "2026-01-20", "time": "14-30-00", "dimensions": "512x512"}
result = build_metadata_section(info, include_date=False, include_time=False)
self.assertNotIn("Date", result)
self.assertNotIn("Time", result)
self.assertIn("Dimensions", result)
def test_trailing_newline(self):
"""Test that section ends with newline."""
result = build_metadata_section({"date": "2026-01-20"})
self.assertTrue(result.endswith("\n"))
class TestBuildPromptSection(unittest.TestCase):
"""Test build_prompt_section function."""
def test_no_prompts(self):
"""Test with no prompts returns empty string."""
result = build_prompt_section(None, None)
self.assertEqual(result, "")
def test_empty_prompts(self):
"""Test with empty prompts returns empty string."""
result = build_prompt_section("", "")
self.assertEqual(result, "")
def test_whitespace_prompts(self):
"""Test with whitespace-only prompts returns empty string."""
result = build_prompt_section(" ", " \n ")
self.assertEqual(result, "")
def test_positive_only(self):
"""Test with only positive prompt."""
result = build_prompt_section("a beautiful sunset", None)
self.assertIn("**Positive:**", result)
self.assertIn("a beautiful sunset", result)
self.assertNotIn("**Negative:**", result)
def test_negative_only(self):
"""Test with only negative prompt."""
result = build_prompt_section(None, "blurry, low quality")
self.assertIn("**Negative:**", result)
self.assertIn("blurry, low quality", result)
self.assertNotIn("**Positive:**", result)
def test_both_prompts(self):
"""Test with both prompts."""
result = build_prompt_section("a cat", "dog")
self.assertIn("**Positive:**", result)
self.assertIn("a cat", result)
self.assertIn("**Negative:**", result)
self.assertIn("dog", result)
def test_custom_section_title(self):
"""Test custom section title."""
result = build_prompt_section("test", None, section_title="Custom Prompts")
self.assertIn("**Custom Prompts:**", result)
def test_code_block_formatting(self):
"""Test prompts are wrapped in code blocks."""
result = build_prompt_section("test prompt", None)
self.assertIn("```\ntest prompt\n```", result)
def test_non_string_conversion(self):
"""Test that non-string prompts are converted."""
result = build_prompt_section(12345, None)
self.assertIn("12345", result)
class TestBuildDiscordMessage(unittest.TestCase):
"""Test build_discord_message function."""
def test_empty_message(self):
"""Test building empty message."""
result = build_discord_message()
self.assertEqual(result, "")
def test_base_message_only(self):
"""Test with just base message."""
result = build_discord_message(base_message="Hello!")
self.assertEqual(result, "Hello!")
def test_with_metadata(self):
"""Test with metadata section."""
result = build_discord_message(
base_message="Image generated",
metadata_section="\n**Info:** test"
)
self.assertIn("Image generated", result)
self.assertIn("**Info:** test", result)
def test_with_all_sections(self):
"""Test with all sections."""
result = build_discord_message(
base_message="Base",
metadata_section="\nMeta",
prompt_section="\nPrompt",
additional_sections=["\nExtra1", "\nExtra2"]
)
self.assertIn("Base", result)
self.assertIn("Meta", result)
self.assertIn("Prompt", result)
self.assertIn("Extra1", result)
self.assertIn("Extra2", result)
def test_truncation(self):
"""Test message truncation at max length."""
long_message = "x" * 2500
result = build_discord_message(base_message=long_message, max_length=2000)
self.assertLessEqual(len(result), 2000)
self.assertIn("[Message truncated]", result)
def test_no_truncation_under_limit(self):
"""Test message not truncated when under limit."""
message = "x" * 100
result = build_discord_message(base_message=message)
self.assertNotIn("truncated", result)
class TestValidateMessageContent(unittest.TestCase):
"""Test validate_message_content function."""
def test_empty_message(self):
"""Test empty message is valid."""
is_valid, msg = validate_message_content("")
self.assertTrue(is_valid)
self.assertIn("Empty message", msg)
def test_normal_message(self):
"""Test normal message is valid."""
is_valid, msg = validate_message_content("Hello world")
self.assertTrue(is_valid)
def test_too_long_message(self):
"""Test message over 2000 chars is invalid."""
is_valid, msg = validate_message_content("x" * 2001)
self.assertFalse(is_valid)
self.assertIn("2000 character limit", msg)
def test_message_with_prompts_section(self):
"""Test message with Generation Prompts section."""
message = "Test\n**Generation Prompts:**\nContent"
is_valid, msg = validate_message_content(message)
self.assertTrue(is_valid)
self.assertNotIn("WARNING", msg)
def test_message_without_prompts_section(self):
"""Test message without Generation Prompts section shows warning."""
is_valid, msg = validate_message_content("Test message")
self.assertTrue(is_valid)
self.assertIn("WARNING", msg)
class TestFormatFileSize(unittest.TestCase):
"""Test format_file_size function."""
def test_bytes(self):
"""Test formatting bytes."""
self.assertEqual(format_file_size(500), "500 bytes")
def test_kilobytes(self):
"""Test formatting kilobytes."""
self.assertEqual(format_file_size(2048), "2.0 KB")
def test_megabytes(self):
"""Test formatting megabytes."""
self.assertEqual(format_file_size(5 * 1024 * 1024), "5.0 MB")
def test_gigabytes(self):
"""Test formatting gigabytes."""
self.assertEqual(format_file_size(2 * 1024 * 1024 * 1024), "2.00 GB")
def test_zero(self):
"""Test formatting zero bytes."""
self.assertEqual(format_file_size(0), "0 bytes")
class TestFormatFileInfo(unittest.TestCase):
"""Test format_file_info function."""
def test_basic_info(self):
"""Test basic file info formatting."""
result = format_file_info("image.png", 1024)
self.assertIn("image.png", result)
self.assertIn("1.0 KB", result)
def test_with_mime_type(self):
"""Test file info with MIME type."""
result = format_file_info("video.mp4", 1024 * 1024, "video/mp4")
self.assertIn("video.mp4", result)
self.assertIn("1.0 MB", result)
self.assertIn("[video/mp4]", result)
if __name__ == "__main__":
unittest.main()