test_molecule_image.py

#!/usr/bin/env python3
"""
Test script for molecular structure image generation.
"""

from rdkit import Chem
from rdkit.Chem import Draw

def generate_molecule_image(smiles):
    """Generate a molecular structure image from SMILES string."""
    mol = Chem.MolFromSmiles(smiles)
    if not mol:
        return None
    
    # Create a high-quality image
    img = Draw.MolToImage(mol, size=(300, 300), kekulize=True)
    return img

def test_molecule_images():
    """Test molecular structure generation for various molecules."""
    
    test_molecules = [
        ("CC", "Ethane"),
        ("C[C@H](N)C(=O)O", "Alanine"),
        ("CC(=O)OC1=CC=CC=C1C(=O)O", "Aspirin"),
        ("CC1=CC=C(C=C1)C(C)NC(=O)C2=CC=C(C=C2)Cl", "Ibuprofen"),
        ("invalid_smiles", "Invalid SMILES")
    ]
    
    print("🧬 Testing Molecular Structure Generation")
    print("=" * 50)
    
    for smiles, name in test_molecules:
        print(f"\n📊 Testing: {name}")
        print(f"SMILES: {smiles}")
        
        img = generate_molecule_image(smiles)
        if img:
            print(f"✅ Successfully generated {img.size[0]}x{img.size[1]} image")
            # Save the image for verification
            img.save(f"test_{name.replace(' ', '_').lower()}.png")
            print(f"💾 Saved as test_{name.replace(' ', '_').lower()}.png")
        else:
            print("❌ Failed to generate image")
        print("-" * 30)

if __name__ == "__main__":
    test_molecule_images()