diff --git a/Dockerfile b/Dockerfile
index 9aec71981..2255a068f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -34,13 +34,24 @@ RUN make install
 
 ENV HOME="/home/ubuntu"
 ENV JASMINEGRAPH_HOME="/home/ubuntu/software/jasminegraph"
+
 RUN mkdir /home/ubuntu/software/jasminegraph
 WORKDIR /home/ubuntu/software/jasminegraph
+
 RUN pip install tensorflow==2.5.3
 RUN pip install -U scikit-learn
-COPY . .
+COPY ./GraphSAGE ./GraphSAGE
 RUN pip install -r ./GraphSAGE/requirements
 
+COPY ./conf ./conf
+COPY ./src ./src
+COPY ./src_python ./src_python
+COPY ./build.sh ./build.sh
+COPY ./CMakeLists.txt ./CMakeLists.txt
+COPY ./main.cpp ./main.cpp
+COPY ./main.h ./main.h
+COPY ./run-docker.sh ./run-docker.sh
+
 RUN sh build.sh
 ENTRYPOINT ["/home/ubuntu/software/jasminegraph/run-docker.sh"]
 CMD ["bash"]
\ No newline at end of file
diff --git a/test-docker.sh b/test-docker.sh
index edd0c0597..7f36a8b33 100755
--- a/test-docker.sh
+++ b/test-docker.sh
@@ -28,7 +28,7 @@ while ! nc -zvn 127.0.0.1 7777 &>/dev/null; do
     sleep .2
 done
 
-timeout 1800 python3 -u tests/integration/test.py |& tee "logs/${run_id}_test.txt"
+timeout 600 python3 -u tests/integration/test.py |& tee "logs/${run_id}_test.txt"
 exit_code="${PIPESTATUS[0]}"
 rm -rf tests/integration/env
 stop_and_remove_containers
diff --git a/tests/integration/test.py b/tests/integration/test.py
index 8fd77ea62..1a79803ef 100644
--- a/tests/integration/test.py
+++ b/tests/integration/test.py
@@ -32,24 +32,26 @@
 DONE=b'done'
 LINE_END=b'\r\n'
 
-def recv_all(sock, n_bytes):
+def expect_response(sock, expected):
+    global passedAll
     buffer = bytearray()
     read = 0
-    while read < n_bytes:
-        received = sock.recv(n_bytes - read)
+    expected_len = len(expected)
+    while read < expected_len:
+        received = sock.recv(expected_len - read)
+        received_len = len(received)
         if received:
-            read += len(received)
+            if received != expected[read:read + received_len]:
+                buffer.extend(received)
+                data = bytes(buffer)
+                logging.warning(f'Output mismatch\nexpected : {expected}\nreceived : {data}')
+                passedAll = False
+                return False
+            read += received_len
             buffer.extend(received)
-    return bytes(buffer)
-
-def expect_response(sock, expected):
-    global passedAll
-    data = recv_all(sock, len(expected))
+    data = bytes(buffer)
     print(data.decode('utf-8'), end='')
-    if data != expected:
-        logging.warning(f'Output mismatch\nexpected : {expected}\nreceived : {data}')
-        passedAll = False
-        return False
+    assert data == expected
     return True
 
 def send_and_expect_response(sock, testName, send, expected, exitOnFail=False):
@@ -60,7 +62,7 @@ def send_and_expect_response(sock, testName, send, expected, exitOnFail=False):
         failedTests.append(testName)
         if exitOnFail:
             print()
-            logging.fatal('Failed some tests,', file=sys.stderr)
+            logging.fatal('Failed some tests,')
             print(*failedTests, sep='\n', file=sys.stderr)
             sys.exit(1)