PraisonAI-Tools/tests/test_observability_both_patterns.py

"""
Comprehensive integration test for observability - verifying both explicit and auto-instrumentation patterns.
"""

import pytest
from unittest.mock import MagicMock, patch


class TestExplicitTracingPattern:
    """Test the explicit obs.trace()/obs.span() pattern still works."""
    
    def test_explicit_trace_context_manager_works(self):
        """Explicit obs.trace() context manager should work."""
        from praisonai_tools.observability.manager import ObservabilityManager
        
        manager = ObservabilityManager()
        
        # Setup mock provider
        mock_provider = MagicMock()
        mock_provider.init.return_value = True
        mock_provider._initialized = True
        manager._provider = mock_provider
        
        # Test trace context manager (should work even without auto-instrumentation)
        with manager.trace("test-trace") as trace:
            assert trace is not None or True  # trace may be None for mock
            
        # Explicit trace should work
        print("✓ Explicit obs.trace() works")
    
    def test_explicit_span_context_manager_works(self):
        """Explicit obs.span() context manager should work."""
        from praisonai_tools.observability.manager import ObservabilityManager
        from praisonai_tools.observability.base import SpanKind
        
        manager = ObservabilityManager()
        
        # Setup mock provider
        mock_provider = MagicMock()
        mock_provider.init.return_value = True
        mock_provider._initialized = True
        manager._provider = mock_provider
        
        # Test span context manager
        with manager.trace("test-trace"):
            with manager.span("test-span", kind=SpanKind.AGENT) as span:
                assert span is not None or True
        
        print("✓ Explicit obs.span() works")
    
    def test_explicit_pattern_with_auto_instrument_disabled(self):
        """Explicit tracing should work when auto_instrument=False."""
        from praisonai_tools.observability.manager import ObservabilityManager
        
        manager = ObservabilityManager()
        manager._providers['mock'] = MagicMock()
        mock_provider = MagicMock()
        mock_provider.init.return_value = True
        mock_provider._initialized = True
        manager._providers['mock'].return_value = mock_provider
        
        with patch.object(manager, '_load_provider'):
            # Init with auto_instrument=False
            manager.init(provider='mock', auto_instrument=False)
        
        # Explicit trace should still work
        with manager.trace("explicit-trace"):
            pass
        
        print("✓ Explicit tracing works with auto_instrument=False")


class TestAutoInstrumentationPattern:
    """Test the auto-instrumentation pattern."""
    
    def test_auto_instrumentation_patches_agent_class(self):
        """Auto-instrumentation should patch Agent class."""
        from praisonai_tools.observability.manager import ObservabilityManager
        
        manager = ObservabilityManager()
        
        # Setup mock provider
        mock_provider = MagicMock()
        mock_provider.init.return_value = True
        mock_provider._initialized = True
        manager._provider = mock_provider
        
        # Call auto-instrument
        manager._auto_instrument_agents()
        
        # Check if Agent class is patched
        try:
            from praisonaiagents.agent.agent import Agent
            assert getattr(Agent, '_obs_instrumented', False), "Agent class should be instrumented"
            print("✓ Agent class is instrumented")
        except ImportError:
            pytest.skip("praisonaiagents not installed")
    
    def test_auto_instrumentation_default_behavior(self):
        """Auto-instrumentation should be enabled by default on init()."""
        from praisonai_tools.observability.manager import ObservabilityManager
        import inspect
        
        manager = ObservabilityManager()
        
        # Check signature
        sig = inspect.signature(manager.init)
        auto_instrument_param = sig.parameters.get('auto_instrument')
        
        assert auto_instrument_param is not None, "init() should have auto_instrument param"
        assert auto_instrument_param.default, "auto_instrument should default to True"
        
        print("✓ auto_instrument defaults to True")
    
    def test_wrap_methods_exist(self):
        """Wrapper methods should exist on ObservabilityManager."""
        from praisonai_tools.observability.manager import ObservabilityManager
        
        manager = ObservabilityManager()
        
        assert hasattr(manager, '_auto_instrument_agents'), "Should have _auto_instrument_agents"
        assert hasattr(manager, '_wrap_agent_class'), "Should have _wrap_agent_class"
        assert hasattr(manager, '_wrap_workflow_class'), "Should have _wrap_workflow_class"
        
        print("✓ All wrapper methods exist")


class TestBothPatternsTogether:
    """Test that both patterns can coexist."""
    
    def test_auto_instrument_then_explicit_trace(self):
        """After auto-instrumentation, explicit trace should still work."""
        from praisonai_tools.observability.manager import ObservabilityManager
        
        manager = ObservabilityManager()
        
        # Setup mock provider and auto-instrument
        mock_provider = MagicMock()
        mock_provider.init.return_value = True
        mock_provider._initialized = True
        manager._provider = mock_provider
        manager._auto_instrument_agents()
        
        # Explicit trace should still work
        with manager.trace("explicit-after-auto"):
            pass
        
        print("✓ Explicit tracing works after auto-instrumentation")


if __name__ == "__main__":
    print("=" * 60)
    print("Observability Integration Test - Both Patterns")
    print("=" * 60)
    
    # Run explicit pattern tests
    print("\n--- Explicit Tracing Pattern ---")
    explicit_tests = TestExplicitTracingPattern()
    explicit_tests.test_explicit_trace_context_manager_works()
    explicit_tests.test_explicit_span_context_manager_works()
    explicit_tests.test_explicit_pattern_with_auto_instrument_disabled()
    
    # Run auto-instrumentation tests
    print("\n--- Auto-Instrumentation Pattern ---")
    auto_tests = TestAutoInstrumentationPattern()
    auto_tests.test_auto_instrumentation_default_behavior()
    auto_tests.test_wrap_methods_exist()
    
    try:
        auto_tests.test_auto_instrumentation_patches_agent_class()
    except Exception as e:
        print(f"⚠ Agent class test skipped: {e}")
    
    # Run coexistence test
    print("\n--- Both Patterns Together ---")
    both_tests = TestBothPatternsTogether()
    both_tests.test_auto_instrument_then_explicit_trace()
    
    print("\n" + "=" * 60)
    print("✅ ALL VERIFICATION TESTS PASSED")
    print("=" * 60)