diff --git a/examples/ignore_websocket.py b/examples/ignore_websocket.py
index 1ee81d38e..48093951e 100644
--- a/examples/ignore_websocket.py
+++ b/examples/ignore_websocket.py
@@ -14,14 +14,21 @@ from libmproxy.protocol.tcp import TCPHandler
 from libmproxy.protocol import KILL
 from libmproxy.script import concurrent
 
-HTTPRequest._headers_to_strip_off.remove("Connection")
-HTTPRequest._headers_to_strip_off.remove("Upgrade")
+
+def start(context, argv):
+    HTTPRequest._headers_to_strip_off.remove("Connection")
+    HTTPRequest._headers_to_strip_off.remove("Upgrade")
+
+
+def done(context):
+    HTTPRequest._headers_to_strip_off.append("Connection")
+    HTTPRequest._headers_to_strip_off.append("Upgrade")
 
 @concurrent
 def response(context, flow):
-	if flow.response.headers.get_first("Connection", None) == "Upgrade":
-		# We need to send the response manually now...
-		flow.client_conn.send(flow.response.assemble())
-		# ...and then delegate to tcp passthrough.
-		TCPHandler(flow.live.c, log=False).handle_messages()
-		flow.reply(KILL)
\ No newline at end of file
+    if flow.response.headers.get_first("Connection", None) == "Upgrade":
+        # We need to send the response manually now...
+        flow.client_conn.send(flow.response.assemble())
+        # ...and then delegate to tcp passthrough.
+        TCPHandler(flow.live.c, log=False).handle_messages()
+        flow.reply(KILL)
\ No newline at end of file
diff --git a/test/test_examples.py b/test/test_examples.py
index a5a212cd1..deb97b493 100644
--- a/test/test_examples.py
+++ b/test/test_examples.py
@@ -16,4 +16,5 @@ def test_load_scripts():
             f += " foo"  # one argument required
         if "modify_response_body" in f:
             f += " foo bar"  # two arguments required
-        script.Script(f, tmaster)  # Loads the script file.
\ No newline at end of file
+        s = script.Script(f, tmaster)  # Loads the script file.
+        s.unload()
\ No newline at end of file