From ec01660c39f31e5ab4daa3dd02d261f2918bb85a Mon Sep 17 00:00:00 2001
From: Camil Staps
Date: Thu, 17 Mar 2016 08:38:43 +0100
Subject: Project 1: persistent connections in test

---
 project1/proj1_s4498062/webtests.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

(limited to 'project1/proj1_s4498062/webtests.py')

diff --git a/project1/proj1_s4498062/webtests.py b/project1/proj1_s4498062/webtests.py
index b0ff889..3827144 100644
--- a/project1/proj1_s4498062/webtests.py
+++ b/project1/proj1_s4498062/webtests.py
@@ -16,8 +16,8 @@ class TestGetRequests(unittest.TestCase):
 
     def setUp(self):
         """Prepare for testing"""
-        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.client_socket.connect(("localhost", portnr))
+        self.client_skt = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.client_skt.connect(("localhost", portnr))
         self.parser = webhttp.parser.ResponseParser()
         self.default_headers = [
                 ('Host', 'localhost:%d' % portnr),
@@ -28,20 +28,29 @@ class TestGetRequests(unittest.TestCase):
 
     def tearDown(self):
         """Clean up after testing"""
-        self.client_socket.shutdown(socket.SHUT_RDWR)
-        self.client_socket.close()
+        try:
+            self.client_skt.shutdown(socket.SHUT_RDWR)
+            self.client_skt.close()
+        except:
+            pass
 
-    def request(self, method, uri, headers):
+    def request(self, method, uri, headers, persistent=False):
         request = webhttp.message.Request()
         request.method = method
         request.uri = uri
         for name, value in headers:
             request.set_header(name, value)
 
-        self.client_socket.send(str(request))
+        self.client_skt.send(str(request))
 
-        message = self.client_socket.recv(1024)
+        message = self.client_skt.recv(1024)
         response = self.parser.parse_response(message)
+
+        if not persistent:
+            self.client_skt.shutdown(socket.SHUT_RDWR)
+            self.client_skt.close()
+            self.client_skt = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            self.client_skt.connect(("localhost", portnr))
         
         return response
 
-- 
cgit v1.2.3