File size: 6,792 Bytes
8ef403e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
#!/usr/bin/env python3

"""

Simple test script for AbMelt structure generation.

Quick test to verify structure generation functionality works.

"""

import os
import sys
import logging
from pathlib import Path

# Add src to path for imports
sys.path.append(str(Path(__file__).parent / "src"))

# Setup basic logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
logger = logging.getLogger(__name__)

def test_imports():
    """Test that all required modules can be imported."""
    logger.info("Testing imports...")
    
    try:
        from structure_prep import (
            prepare_structure, 
            generate_structure_from_sequences,
            validate_structure,
            get_chain_sequences
        )
        logger.info("βœ“ Successfully imported structure_prep modules")
        return True
    except ImportError as e:
        logger.error(f"βœ— Failed to import structure_prep modules: {e}")
        return False

def test_immune_builder():
    """Test ImmuneBuilder functionality."""
    logger.info("Testing ImmuneBuilder...")
    
    try:
        from structure_prep import generate_structure_from_sequences
        
        # Test sequences (shortened for testing)
        heavy_chain = "QVQLVQSGAEVKKPGASVKVSCKASGYTFTSYWMHWVKQRPGQGLEWIGYINPSRGYTNYNQKFKDKATITADESTSTTAYMELSSLRSEDTAVYYCARGGYSSGYYFDYWGQGTLVTVSS"
        light_chain = "DIQMTQSPSSLSASVGDRVTITCRASQDISNYLNWFQQKPGKAPKLLIYYATSLADGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQGNTFPWTFGQGTKVEIKR"
        
        # Generate structure
        output_file = "test_antibody.pdb"
        generated_file = generate_structure_from_sequences(
            heavy_chain=heavy_chain,
            light_chain=light_chain,
            output_file=output_file
        )
        
        # Check if file was created
        if Path(generated_file).exists():
            logger.info(f"βœ“ ImmuneBuilder generated structure: {generated_file}")
            
            # Test validation
            is_valid = validate_structure(generated_file)
            if is_valid:
                logger.info("βœ“ Structure validation passed")
            else:
                logger.warning("⚠ Structure validation failed")
            
            # Test sequence extraction
            chains = get_chain_sequences(generated_file)
            if chains:
                logger.info(f"βœ“ Extracted chains: {list(chains.keys())}")
            else:
                logger.warning("⚠ Failed to extract chain sequences")
            
            # Cleanup
            Path(generated_file).unlink()
            logger.info("βœ“ Cleaned up test file")
            
            return True
        else:
            logger.error(f"βœ— ImmuneBuilder failed to generate structure")
            return False
            
    except Exception as e:
        logger.error(f"βœ— ImmuneBuilder test failed: {e}")
        return False

def test_prepare_structure():
    """Test the main prepare_structure function."""
    logger.info("Testing prepare_structure function...")
    
    try:
        from structure_prep import prepare_structure
        
        # Test configuration
        config = {
            "paths": {
                "temp_dir": "test_temp",
                "output_dir": "test_output",
                "log_dir": "test_logs"
            }
        }
        
        # Create test directories
        for path in config["paths"].values():
            Path(path).mkdir(parents=True, exist_ok=True)
        
        # Test antibody data
        antibody = {
            "name": "test_antibody",
            "heavy_chain": "QVQLVQSGAEVKKPGASVKVSCKASGYTFTSYWMHWVKQRPGQGLEWIGYINPSRGYTNYNQKFKDKATITADESTSTTAYMELSSLRSEDTAVYYCARGGYSSGYYFDYWGQGTLVTVSS",
            "light_chain": "DIQMTQSPSSLSASVGDRVTITCRASQDISNYLNWFQQKPGKAPKLLIYYATSLADGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQGNTFPWTFGQGTKVEIKR",
            "type": "sequences"
        }
        
        # Run prepare_structure
        structure_files = prepare_structure(antibody, config)
        
        # Verify results
        required_keys = ["pdb_file", "work_dir", "chains"]
        for key in required_keys:
            if key not in structure_files:
                logger.error(f"βœ— Missing required key: {key}")
                return False
        
        # Check if files exist
        pdb_file = Path(structure_files["pdb_file"])
        work_dir = Path(structure_files["work_dir"])
        
        if not pdb_file.exists():
            logger.error(f"βœ— PDB file does not exist: {pdb_file}")
            return False
        
        if not work_dir.exists():
            logger.error(f"βœ— Work directory does not exist: {work_dir}")
            return False
        
        logger.info(f"βœ“ prepare_structure successful")
        logger.info(f"  PDB file: {pdb_file}")
        logger.info(f"  Work dir: {work_dir}")
        logger.info(f"  Chains: {list(structure_files['chains'].keys())}")
        
        # Cleanup
        import shutil
        shutil.rmtree("test_temp")
        shutil.rmtree("test_output")
        shutil.rmtree("test_logs")
        logger.info("βœ“ Cleaned up test directories")
        
        return True
        
    except Exception as e:
        logger.error(f"βœ— prepare_structure test failed: {e}")
        return False

def main():
    """Run all tests."""
    logger.info("=" * 50)
    logger.info("ABMELT STRUCTURE GENERATION TEST")
    logger.info("=" * 50)
    
    tests = [
        ("Import Test", test_imports),
        ("ImmuneBuilder Test", test_immune_builder),
        ("Prepare Structure Test", test_prepare_structure)
    ]
    
    results = []
    
    for test_name, test_func in tests:
        logger.info(f"\n{test_name}:")
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            logger.error(f"βœ— {test_name} failed with exception: {e}")
            results.append((test_name, False))
    
    # Summary
    logger.info("\n" + "=" * 50)
    logger.info("TEST SUMMARY")
    logger.info("=" * 50)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = "βœ“ PASS" if result else "βœ— FAIL"
        logger.info(f"{test_name}: {status}")
    
    logger.info(f"\nOverall: {passed}/{total} tests passed")
    
    if passed == total:
        logger.info("πŸŽ‰ All tests passed!")
        return 0
    else:
        logger.error("❌ Some tests failed")
        return 1

if __name__ == "__main__":
    sys.exit(main())