File size: 3,305 Bytes
a2e1879
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#!/usr/bin/env python3
"""
Simple test script to verify attachment guardrails functionality
"""

import os
import sys

# Add current directory to path
sys.path.insert(0, '.')

def test_attachment_guardrails():
    """Test the attachment guardrail system"""
    try:
        # Test basic imports
        print("Testing imports...")
        from guardrails.attachments.base import AttachmentGuardrailManager
        from guardrails.attachments.txt_guardrail import TxtGuardrail
        from guardrails.attachments.pdf_guardrail import PdfGuardrail
        from guardrails.attachments.docx_guardrail import DocxGuardrail
        import config
        
        print("✅ All imports successful")
        
        # Test configuration
        print("\nTesting configuration...")
        print(f"Attachment config: {config.ATTACHMENT_GUARDRAILS_CONFIG}")
        
        # Test guardrail manager initialization
        print("\nTesting guardrail manager...")
        manager = AttachmentGuardrailManager(config.ATTACHMENT_GUARDRAILS_CONFIG)
        
        print(f"Supported extensions: {manager.get_supported_extensions()}")
        print(f"Guardrail info: {manager.get_guardrail_info()}")
        
        # Test with a simple text file
        print("\nTesting with sample text...")
        sample_text = "Hello world, this is a test file."
        sample_bytes = sample_text.encode('utf-8')
        
        is_safe, analysis = manager.process_attachment("test.txt", sample_bytes)
        
        print(f"Text file - Is safe: {is_safe}")
        print(f"Text file - Analysis summary: chunks={analysis.get('chunks_analyzed', 0)}, confidence_threshold={analysis.get('confidence_threshold', 0)}")
        
        # Test PDF file handling (without actual PDF content)
        print("\nTesting PDF file handling...")
        
        pdf_result = manager.process_attachment("test.pdf", b"dummy content")
        print(f"PDF file - Is safe: {pdf_result[0]}")
        print(f"PDF file - Error (expected for dummy content): {pdf_result[1].get('error', 'No error')}")
        print(f"PDF file - Guardrail used: {pdf_result[1].get('guardrail_used', 'Unknown')}")
        print(f"PDF file - Confidence threshold: {pdf_result[1].get('confidence_threshold', 'N/A')}")
        
        # Test Word document file handling (without actual DOCX content)
        print("\nTesting Word document file handling...")
        
        docx_result = manager.process_attachment("test.docx", b"dummy content")
        print(f"Word file - Is safe: {docx_result[0]}")
        print(f"Word file - Error (expected for dummy content): {docx_result[1].get('error', 'No error')}")
        print(f"Word file - Guardrail used: {docx_result[1].get('guardrail_used', 'Unknown')}")
        print(f"Word file - Confidence threshold: {docx_result[1].get('confidence_threshold', 'N/A')}")
        
        return True
        
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("🧪 Testing Attachment Guardrails System")
    print("=" * 50)
    
    success = test_attachment_guardrails()
    
    if success:
        print("\n✅ All tests passed!")
    else:
        print("\n❌ Tests failed!")
        sys.exit(1)