import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;

import org.junit.Test;
import org.python.core.PyModule;
import org.python.core.PyObject;
import org.python.core.imp;
import org.python.util.PythonInterpreter;
import org.python.util.PythonObjectInputStream;

public class SerializationTests {
       private PythonInterpreter interp;

       public void init() {
               interp = new PythonInterpreter();
       }

       public void setupMain() {
               final PyModule mod = imp.addModule("__main__");
               interp.setLocals(mod.__dict__);
       }

       public void createObject() {
               interp.exec("from java.lang import Object");
               interp.exec("from java.io import Serializable");
               interp.exec("class Test(Serializable):\n\tdef
__init__(self):\n\t\tObject.__init__(self)\n");
               interp.exec("x = Test()");
       }

       public byte[] serialize() throws IOException {
               final PyObject x = interp.get("x");
               final ByteArrayOutputStream os = new ByteArrayOutputStream();
               new ObjectOutputStream(os).writeObject(x);
               return os.toByteArray();
       }

       public void deserialize(final byte[] b) throws IOException,
ClassNotFoundException {
               final ByteArrayInputStream is = new ByteArrayInputStream(b);
               new PythonObjectInputStream(is).readObject();
       }

       public void testSerialization() throws IOException, ClassNotFoundException
{
               deserialize(serialize());
       }

       @Test
       public void testDirect() throws IOException, ClassNotFoundException {
               init();
               createObject();
               testSerialization();
       }

       @Test
       public void testJython() {
               init();
               createObject();
               interp.set("t", this);
               interp.exec("t.testSerialization()");
       }

       @Test
       public void testDirectWithMain() throws IOException, ClassNotFoundException
{
               init();
               setupMain();
               createObject();
               testSerialization();
       }

       @Test
       public void testJythonWithMain() {
               init();
               setupMain();
               createObject();
               interp.set("t", this);
               interp.exec("t.testSerialization()");
       }
}