Add tests for binary copy function

backend/entityservice/database/selections.py

@@ -329,18 +329,19 @@ def copy_binary_column_from_select_query(cur, select_query, stored_binary_size=1
 
     :param select_query: An sql query that select's a single binary column. Include ordering the results.
     :param stored_binary_size: Fixed size of each bytea data.
+    :raises AssertionError if the database implements an unhandled extension or the EOF is corrupt.
     """
 
     copy_to_stream_query = """COPY ({}) TO STDOUT WITH binary""".format(select_query)
     stream = io.BytesIO()
     cur.copy_expert(copy_to_stream_query, stream)
-    # TODO: It may be more efficient to introduce a buffered binary stream instead of getting the entire stream at once
+
     raw_data = stream.getvalue()
 
     # Need to read/remove the Postgres Binary Header, Trailer, and the per tuple info
     # https://www.postgresql.org/docs/current/sql-copy.html
     _ignored_header = raw_data[:15]
-    header_extension = raw_data[16:20]
+    header_extension = raw_data[15:19]
     assert header_extension == b'\x00\x00\x00\x00', "Need to implement skipping postgres binary header extension"
     binary_trailer = raw_data[-2:]
     assert binary_trailer == b'\xff\xff', "Corrupt COPY of binary data from postgres"

backend/entityservice/integrationtests/dbtests/conftest.py

@@ -0,0 +1,57 @@
+import pytest
+import psycopg2
+
+from entityservice.settings import Config as config
+
+
+@pytest.fixture
+def conn():
+    db = config.DATABASE
+    host = config.DATABASE_SERVER
+    user = config.DATABASE_USER
+    password = config.DATABASE_PASSWORD
+    conn = psycopg2.connect(host=host, dbname=db, user=user, password=password)
+    yield conn
+    conn.close()
+
+@pytest.fixture
+def cur(conn):
+    return conn.cursor()
+
+
+
+
+@pytest.fixture()
+def prepopulated_binary_test_data(conn, cur, num_bytes=4, num_rows=100):
+    creation_sql = """
+        DROP TABLE IF EXISTS binary_test;
+        CREATE TABLE binary_test
+            (
+                id          integer not null,
+                encoding    bytea   not null
+            );"""
+    cur.execute(creation_sql)
+    conn.commit()
+
+    # Add data using execute_values
+    data = [(i, bytes([i % 128] * num_bytes)) for i in range(num_rows)]
+    psycopg2.extras.execute_values(cur, """
+    INSERT INTO binary_test (id, encoding) VALUES %s
+    """, data)
+
+    conn.commit()
+
+    # quick check data is there
+    cur.execute("select count(*) from binary_test")
+    res = cur.fetchone()[0]
+    assert res == num_rows
+
+    cur.execute("select encoding from binary_test where id = 1")
+    assert bytes(cur.fetchone()[0]) == data[1][1]
+
+    yield data
+
+    # delete test table
+    deletion_sql = "drop table if exists binary_test cascade;"
+    cur.execute(deletion_sql)
+    conn.commit()

backend/entityservice/integrationtests/dbtests/test_insertions.py

@@ -4,7 +4,7 @@
 from pytest import raises
 
 from entityservice.database import insert_dataprovider, insert_encodings_into_blocks, insert_blocking_metadata, \
-    get_project, get_encodingblock_ids, get_block_metadata, get_chunk_of_encodings
+    get_project, get_encodingblock_ids, get_block_metadata, get_chunk_of_encodings, copy_binary_column_from_select_query
 
 from entityservice.integrationtests.dbtests import _get_conn_and_cursor
 from entityservice.models import Project
@@ -13,8 +13,29 @@
 from entityservice.utils import generate_code
 
 
+class TestBinaryCopy:
+
+    def test_copy_binary_column_from_select_query(self, conn, cur, prepopulated_binary_test_data):
+        query = "select encoding from binary_test where id >= 10 and id < 20"
+        res = list(copy_binary_column_from_select_query(cur, query, stored_binary_size=4))
+        assert len(res) == 10
+        for (i, original), stored in zip(prepopulated_binary_test_data[10:20], res):
+            assert original == stored
+
+    def test_copy_binary_column_from_select_query_empty(self, conn, cur, prepopulated_binary_test_data):
+        query = "select encoding from binary_test where id < 0"
+        res = list(copy_binary_column_from_select_query(cur, query, stored_binary_size=4))
+        assert len(res) == 0
+
+
 class TestInsertions:
 
+    def _create_project(self):
+        project = Project('groups', {}, name='', notes='', parties=2, uses_blocking=False)
+        conn, cur = _get_conn_and_cursor()
+        dp_ids = project.save(conn)
+        return project, dp_ids
+
     def _create_project_and_dp(self):
         project, dp_ids = self._create_project()
         dp_id = dp_ids[0]
@@ -28,11 +49,6 @@ def _create_project_and_dp(self):
         assert len(dp_auth_token) == 48
         return project.project_id, project.result_token, dp_id, dp_auth_token
 
-    def _create_project(self):
-        project = Project('groups', {}, name='', notes='', parties=2, uses_blocking=False)
-        conn, cur = _get_conn_and_cursor()
-        dp_ids = project.save(conn)
-        return project, dp_ids
 
     def test_insert_project(self):
         before = datetime.datetime.now()