package org.postgresql.test.jdbc2;

import org.postgresql.test.TestUtil;
import junit.framework.TestCase;
import java.io.*;
import java.sql.*;

import org.postgresql.copy.*;


public class CopyTest extends TestCase
{
	private Connection con;
	private CopyManager copyMgr;

	protected void setUp() throws Exception
	{
		con = TestUtil.openDB();
		TestUtil.createTable(con, "testcopy", "col1 text, col2 text");
		copyMgr = ((org.postgresql.PGConnection)con).getCopyAPI();
	}

	protected void tearDown() throws Exception
	{
		TestUtil.dropTable(con, "testcopy");
		TestUtil.closeDB(con);
	}


	/*
	 * Tests the copyIn functionality by copying in an array of bytes, then checking their integrity in the table
	 */
	public void testCopyIn()
	{
		try
		{
			java.sql.Statement st = con.createStatement();
			byte[] dummy_array = new byte[] { (byte)'T', (byte)'o', (byte)'m', (byte)'\t', (byte)'1', (byte)'2', (byte)'\n',	//Tom 12
											  (byte)'B', (byte)'o', (byte)'b', (byte)'\t', (byte)'2', (byte)'9', (byte)'\n',	//Bob 29
											  (byte)'H', (byte)'a', (byte)'n', (byte)'\t', (byte)'4', (byte)'1', (byte)'\n'    //Han 41
											};

			InputStream in = new ByteArrayInputStream(dummy_array);
			copyMgr.copyIn("testcopy", in);

			java.sql.ResultSet rs = st.executeQuery( "SELECT sum(col2::int) AS sum FROM testcopy" );
			rs.first();
			int sum = rs.getInt("sum");
			assertEquals(sum, 82); // the sum of column col2 should equal 82 (i.e. 12 + 29 + 41)

		}
		catch (Exception ex)
		{
			assertTrue(ex.getMessage(), false);
		}
	}

	/*
	 * Tests the copyOut functionality by copying data from pg_class into a test table,
	 * then it copies that data out and check the number of 'rows' that came out.
	 */
	public void testCopyOut()
	{
		try
		{
			java.sql.Statement st = con.createStatement();
			st.executeUpdate( "INSERT into testcopy SELECT relowner, relpages FROM pg_class LIMIT 20" );

			ByteArrayOutputStream out = new ByteArrayOutputStream();
			java.sql.ResultSet rs = st.executeQuery( "SELECT count(*) AS row_count FROM testcopy" );
			rs.first();
			int row_count = rs.getInt("row_count"); // the number of rows in the table testcopy

			copyMgr.copyOut("testcopy", out);

			ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
			int newline_count = 0;
			for ( int i = in.read(); i != -1; i = in.read() )
			{
				if (i == '\n')
					newline_count++; // count the number of newlines
			}

			// # of newlines should equal original LIMIT #
			assertEquals(row_count, newline_count);
		}
		catch (Exception ex)
		{
			assertTrue(ex.getMessage(), false);
		}
	}








	/*
	 * Test a sequence of copy commands by
	 *
	 */
	public void testCopyOutCopyIn()
	{
		try
		{
			java.sql.Statement st = con.createStatement();

			st.executeUpdate( "INSERT into testcopy SELECT relowner, relpages FROM pg_class LIMIT 20" );

			ByteArrayOutputStream out = new ByteArrayOutputStream();
			java.sql.ResultSet rs = st.executeQuery( "SELECT count(*) AS row_count FROM testcopy" );
			rs.first();
			int row_count = rs.getInt("row_count"); // the number of rows in the table testcopy

			copyMgr.copyOut("testcopy", out);

			ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
			int newline_count = 0;
			for ( int i = in.read(); i != -1; i = in.read() )
			{
				if (i == '\n')
					newline_count++; // count the number of newlines
			}

			// # of newlines should equal original LIMIT #
			assertEquals(row_count, newline_count);

			in = new ByteArrayInputStream(out.toByteArray());
			copyMgr.copyIn("testcopy", in);

			rs = st.executeQuery( "SELECT count(*) AS row_count FROM testcopy" );
			rs.first();
			row_count = rs.getInt("row_count"); // the number of rows in the table testcopy

			assertEquals(row_count , newline_count*2);
		}
		catch (Exception ex)
		{
			assertTrue(ex.getMessage(), false);
		}
	}




	/*
	 * Test a sequence of copy commands by copying in an array of bytes, then copying them out
	 * and comparing the lengths of the two arrays.
	 */
	public void testCopyInCopyOut()
	{
		try
		{
			java.sql.Statement st = con.createStatement();
			byte[] input_array = new byte[] { (byte)'T', (byte)'o', (byte)'m', (byte)'\t', (byte)'1', (byte)'2', (byte)'\n',
											  (byte)'B', (byte)'o', (byte)'b', (byte)'\t', (byte)'2', (byte)'9', (byte)'\n',
											  (byte)'H', (byte)'a', (byte)'n', (byte)'\t', (byte)'4', (byte)'1', (byte)'\n'
											};

			InputStream in = new ByteArrayInputStream(input_array);
			copyMgr.copyIn("testcopy", in);

			java.sql.ResultSet rs = st.executeQuery( "SELECT sum(col2::int) AS sum FROM testcopy" );
			rs.first();
			int sum = rs.getInt("sum");
			assertEquals(sum, 82); // the sum of column 'col2' should equal 82 (i.e. 12 + 29 + 41)

			// copy out the everything from the table we just copy'ed into

			ByteArrayOutputStream out = new ByteArrayOutputStream();
			copyMgr.copyOut("testcopy", out);

			int input_length = input_array.length;
			int output_length = out.size();
			assertEquals(input_length, output_length); // the input and output array should be the same size
		}
		catch (Exception ex)
		{
			assertTrue(ex.getMessage(), false);
		}
	}
}