test_zip.py

import zipfile
import tempfile
import os

def create_xyz_zip(xyz_strings):
    """Create a zip file containing all generated XYZs."""
    if not xyz_strings:
        return None
        
    # Create a temporary file
    tmp_zip = tempfile.NamedTemporaryFile(suffix='.zip', delete=False)
    tmp_zip.close()
    
    with zipfile.ZipFile(tmp_zip.name, 'w') as zf:
        for i, xyz_content in enumerate(xyz_strings):
            # writestr takes (archive_name, data)
            zf.writestr(f"molecule_{i:03d}.xyz", xyz_content)
            
    return tmp_zip.name

# Test usage
long_string = "A" * 1000  # Simulate long xyz content
try:
    zip_path = create_xyz_zip([long_string])
    print(f"SUCCESS: Zip created at {zip_path}")
    # Verify content
    with zipfile.ZipFile(zip_path, 'r') as zf:
        print(f"Files in zip: {zf.namelist()}")
        content = zf.read("molecule_000.xyz").decode('utf-8')
        if content == long_string:
            print("Content matches!")
        else:
            print("Content mismatches!")
    os.remove(zip_path)
except Exception as e:
    print(f"FAIL: {e}")