import pickle
import pytest

try:
    import pycuda
except ImportError:
    ans = raw_input('PyCUDA not found.  Regression tests will take forever.  Do you want to continue? [y/n] ')
    if ans in ['Y', 'y']:
        pass
    else:
        sys.exit()

from pygbe.main import main

@pytest.mark.parametrize('key', ['total_elements',
                                 'E_solv_kJ',
                                 'E_coul_kcal',
                                 'E_coul_kJ',
                                 'E_solv_kcal'])
def test_lysozyme(key):
    results = main(['', '../examples/lys',
                    'log_output=False',
                    'return_results_dict=True'])

    with open('lysozyme.pickle', 'r') as f:
        base_results = pickle.load(f)

    assert base_results[key] == results[key]

@@ -25,51 +25,184 @@
 est%0A
-%0Atry:%0A    import pycuda%0Aexcept ImportError:
+# TODO change this to %60import functools%60 for py3 port%0Aimport functools32 as functools%0A%0Atry:%0A    import pycuda%0Aexcept ImportError:%0A    # TODO change this to %60input%60 for py3 port
 %0A   
@@ -689,124 +689,20 @@
 s = 
-main(%5B'', '../examples/lys',%0A                    'log_output=False',%0A                    'return_results_dict=True'%5D
+get_results(
 )%0A%0A 
@@ -828,8 +828,255 @@
 ts%5Bkey%5D%0A
+%0A@functools.lru_cache(5)%0Adef get_results():%0A    print('Generating results for lysozyme example...')%0A    results = main(%5B'','../examples/lys'%5D,%0A                    log_output=False,%0A                    return_results_dict=True)%0A%0A    return results%0A

fix lysozyme test