test_dpg_ui.py

#!/usr/bin/env python3
"""
Comprehensive test script for the Dear PyGui UI

This script allows testing various components of the Dear PyGui UI
without running the full application.
"""

import os
import sys
import dearpygui.dearpygui as dpg
import argparse
from typing import Optional, Dict, Any, List

# Import dark theme if available
try:
    from dpg_ui.themes.dark_theme import apply_dark_theme
except ImportError:
    # Fallback if theme is not available
    def apply_dark_theme():
        pass

# Add the parent directory to sys.path
parent_dir = os.path.abspath(os.path.dirname(__file__))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# Import state manager
from dpg_ui.core.state_manager import StateManager


class TestApp:
    """Test application for Dear PyGui UI components"""
    
    def __init__(self, test_component: str = "concept"):
        """
        Initialize the test application
        
        Args:
            test_component: The component to test
        """
        self.test_component = test_component
        self.state = StateManager()
        self.window_id = "main_window"
        
    def run(self):
        """Run the test application"""
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title=f"OneTrainer - {self.test_component.capitalize()} Test", 
                           width=1024, height=768)
        dpg.setup_dearpygui()
        
        # Apply dark theme
        apply_dark_theme()
        
        # Create main window
        with dpg.window(tag=self.window_id, label=f"OneTrainer - {self.test_component.capitalize()} Test"):
            # Add the component to test
            self.add_test_component()
        
        # Show the viewport
        dpg.show_viewport()
        dpg.set_primary_window(self.window_id, True)
        
        # Start the main loop
        dpg.start_dearpygui()
        
        # Clean up
        dpg.destroy_context()
        
    def add_test_component(self):
        """Add the component to test"""
        if self.test_component == "concept":
            self.test_concept_tab()
        elif self.test_component == "concept_window":
            self.test_concept_window()
        elif self.test_component == "lora":
            self.test_lora_tab()
        elif self.test_component == "model":
            self.test_model_tab()
        elif self.test_component == "training":
            self.test_training_tab()
        elif self.test_component == "sampling":
            self.test_sampling_tab()
        else:
            dpg.add_text(f"Unknown component: {self.test_component}")
    
    def test_concept_tab(self):
        """Test the concept tab"""
        try:
            from dpg_ui.tabs.concept_tab import ConceptTab
            dpg.add_text("Testing Concept Tab")
            dpg.add_separator()
            concept_tab = ConceptTab(dpg.last_item(), self.state)
        except ImportError as e:
            dpg.add_text(f"Error importing Concept Tab: {e}")
    
    def test_concept_window(self):
        """Test the concept window"""
        try:
            from dpg_ui.windows.concept_window_simple import ConceptWindow
            from modules.util.config.ConceptConfig import ConceptConfig
            
            dpg.add_text("Testing Concept Window")
            dpg.add_separator()
            
            # Create a concept config
            concept = ConceptConfig.default_values()
            concept.name = "Test Concept"
            concept.path = os.path.join(os.getcwd(), "training_concepts")
            
            # Create a button to open the concept window
            def open_concept_window():
                def on_concept_updated(updated_concept):
                    if updated_concept:
                        print(f"Concept updated: {updated_concept.name}")
                        
                concept_window = ConceptWindow(concept, self.state, on_concept_updated)
            
            dpg.add_button(label="Open Concept Window", callback=open_concept_window)
            
        except ImportError as e:
            dpg.add_text(f"Error importing Concept Window: {e}")
    
    def test_lora_tab(self):
        """Test the lora tab"""
        try:
            # Try to import the lora tab
            try:
                from dpg_ui.tabs.lora_tab import LoraTab
                dpg.add_text("Testing LoRA Tab")
                dpg.add_separator()
                lora_tab = LoraTab(dpg.last_item(), self.state)
            except ImportError:
                # Try the simplified version
                from dpg_ui.tabs.lora_tab_simple import LoraTab
                dpg.add_text("Testing Simplified LoRA Tab")
                dpg.add_separator()
                lora_tab = LoraTab(dpg.last_item(), self.state)
        except ImportError as e:
            dpg.add_text(f"Error importing LoRA Tab: {e}")
    
    def test_model_tab(self):
        """Test the model tab"""
        try:
            from dpg_ui.tabs.model_tab import ModelTab
            dpg.add_text("Testing Model Tab")
            dpg.add_separator()
            model_tab = ModelTab(dpg.last_item(), self.state)
        except ImportError as e:
            dpg.add_text(f"Error importing Model Tab: {e}")
    
    def test_training_tab(self):
        """Test the training tab"""
        try:
            from dpg_ui.tabs.training_tab import TrainingTab
            dpg.add_text("Testing Training Tab")
            dpg.add_separator()
            training_tab = TrainingTab(dpg.last_item(), self.state)
        except ImportError as e:
            dpg.add_text(f"Error importing Training Tab: {e}")
    
    def test_sampling_tab(self):
        """Test the sampling tab"""
        try:
            from dpg_ui.tabs.sampling_tab import SamplingTab
            dpg.add_text("Testing Sampling Tab")
            dpg.add_separator()
            sampling_tab = SamplingTab(dpg.last_item(), self.state)
        except ImportError as e:
            dpg.add_text(f"Error importing Sampling Tab: {e}")


def main():
    """Main entry point"""
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Test DPG UI components")
    parser.add_argument("--component", type=str, default="concept",
                      choices=["concept", "concept_window", "lora", "model", 
                              "training", "sampling"],
                      help="Component to test")
    args = parser.parse_args()
    
    # Create and run the test application
    app = TestApp(args.component)
    app.run()


if __name__ == "__main__":
    main()