Merge pull request #1192 from cortesi/testsuite

WIP: Solidify pathod test suite

11 files changed, 152 insertions, 84 deletions

diff --git a/README.rst b/README.rst
index b108d26f..4f84effd 100644
--- a/README.rst
+++ b/README.rst
@@ -3,19 +3,24 @@ mitmproxy
 
 |travis| |coveralls| |latest_release| |python_versions|
 
-This repository contains the **mitmproxy** and **pathod** projects, as well as their shared networking library, **netlib**.
+This repository contains the **mitmproxy** and **pathod** projects, as well as
+their shared networking library, **netlib**.
 
-``mitmproxy`` is an interactive, SSL-capable intercepting proxy with a console interface.
+``mitmproxy`` is an interactive, SSL-capable intercepting proxy with a console
+interface.
 
 ``mitmdump`` is the command-line version of mitmproxy. Think tcpdump for HTTP.
 
-``pathoc`` and ``pathod`` are perverse HTTP client and server applications designed to let you craft almost any conceivable HTTP request, including ones that creatively violate the standards.
+``pathoc`` and ``pathod`` are perverse HTTP client and server applications
+designed to let you craft almost any conceivable HTTP request, including ones
+that creatively violate the standards.
 
 
 Documentation & Help
 --------------------
 
-Documentation, tutorials and precompiled binaries can be found on the mitmproxy and pathod websites.
+Documentation, tutorials and precompiled binaries can be found on the mitmproxy
+and pathod websites.
 
 |mitmproxy_site| |pathod_site|
 
@@ -39,8 +44,8 @@ Hacking
 -------
 
 To get started hacking on mitmproxy, make sure you have Python_ 2.7.x. with
-virtualenv_ installed (you can find installation instructions for virtualenv here_).
-Then do the following:
+virtualenv_ installed (you can find installation instructions for virtualenv
+here_). Then do the following:
 
 .. code-block:: text
 
@@ -49,10 +54,11 @@ Then do the following:
     ./dev.sh
 
 
-The *dev* script will create a virtualenv environment in a directory called "venv",
-and install all mandatory and optional dependencies into it.
-The primary mitmproxy components - mitmproxy, netlib and pathod - are installed as "editable",
-so any changes to the source in the repository will be reflected live in the virtualenv.
+The *dev* script will create a virtualenv environment in a directory called
+"venv", and install all mandatory and optional dependencies into it. The
+primary mitmproxy components - mitmproxy, netlib and pathod - are installed as
+"editable", so any changes to the source in the repository will be reflected
+live in the virtualenv.
 
 To confirm that you're up and running, activate the virtualenv, and run the
 mitmproxy test suite:
@@ -63,9 +69,9 @@ mitmproxy test suite:
     py.test
 
 Note that the main executables for the project - ``mitmdump``, ``mitmproxy``,
-``mitmweb``, ``pathod``, and ``pathoc`` - are all created within the virtualenv. After activating the
-virtualenv, they will be on your $PATH, and you can run them like any other
-command:
+``mitmweb``, ``pathod``, and ``pathoc`` - are all created within the
+virtualenv. After activating the virtualenv, they will be on your $PATH, and
+you can run them like any other command:
 
 .. code-block:: text
 
@@ -92,9 +98,9 @@ suite. The project tries to maintain 100% test coverage.
 Documentation
 -------------
 
-The mitmproxy documentation is build using Sphinx_, which is installed automatically if you set up a development
-environment as described above.
-After installation, you can render the documentation like this:
+The mitmproxy documentation is build using Sphinx_, which is installed
+automatically if you set up a development environment as described above. After
+installation, you can render the documentation like this:
 
 .. code-block:: text
 
diff --git a/netlib/tcp.py b/netlib/tcp.py
index 914aa701..de12102e 100644
--- a/netlib/tcp.py
+++ b/netlib/tcp.py
@@ -6,6 +6,7 @@ import sys
 import threading
 import time
 import traceback
+import contextlib
 
 import binascii
 from six.moves import range
@@ -577,6 +578,12 @@ class _Connection(object):
         return context
 
 
+@contextlib.contextmanager
+def _closer(client):
+    yield
+    client.close()
+
+
 class TCPClient(_Connection):
 
     def __init__(self, address, source_address=None):
@@ -708,6 +715,7 @@ class TCPClient(_Connection):
         self.connection = connection
         self.ip_address = Address(connection.getpeername())
         self._makefile()
+        return _closer(self)
 
     def settimeout(self, n):
         self.connection.settimeout(n)
@@ -833,6 +841,25 @@ class BaseHandler(_Connection):
             return b""
 
 
+class Counter:
+    def __init__(self):
+        self._count = 0
+        self._lock = threading.Lock()
+
+    @property
+    def count(self):
+        with self._lock:
+            return self._count
+
+    def __enter__(self):
+        with self._lock:
+            self._count += 1
+
+    def __exit__(self, *args):
+        with self._lock:
+            self._count -= 1
+
+
 class TCPServer(object):
     request_queue_size = 20
 
@@ -845,15 +872,17 @@ class TCPServer(object):
         self.socket.bind(self.address())
         self.address = Address.wrap(self.socket.getsockname())
         self.socket.listen(self.request_queue_size)
+        self.handler_counter = Counter()
 
     def connection_thread(self, connection, client_address):
-        client_address = Address(client_address)
-        try:
-            self.handle_client_connection(connection, client_address)
-        except:
-            self.handle_error(connection, client_address)
-        finally:
-            close_socket(connection)
+        with self.handler_counter:
+            client_address = Address(client_address)
+            try:
+                self.handle_client_connection(connection, client_address)
+            except:
+                self.handle_error(connection, client_address)
+            finally:
+                close_socket(connection)
 
     def serve_forever(self, poll_interval=0.1):
         self.__is_shut_down.clear()
diff --git a/pathod/pathoc.py b/pathod/pathoc.py
index 2b7d053c..5cfb4591 100644
--- a/pathod/pathoc.py
+++ b/pathod/pathoc.py
@@ -286,7 +286,7 @@ class Pathoc(tcp.TCPClient):
         if self.use_http2 and not self.ssl:
             raise NotImplementedError("HTTP2 without SSL is not supported.")
 
-        tcp.TCPClient.connect(self)
+        ret = tcp.TCPClient.connect(self)
 
         if connect_to:
             self.http_connect(connect_to)
@@ -324,6 +324,7 @@ class Pathoc(tcp.TCPClient):
 
         if self.timeout:
             self.settimeout(self.timeout)
+        return ret
 
     def stop(self):
         if self.ws_framereader:
diff --git a/pathod/pathod.py b/pathod/pathod.py
index 7795df0e..0449c0c1 100644
--- a/pathod/pathod.py
+++ b/pathod/pathod.py
@@ -353,6 +353,8 @@ class Pathod(tcp.TCPServer):
             staticdir=self.staticdir
         )
 
+        self.loglock = threading.Lock()
+
     def check_policy(self, req, settings):
         """
             A policy check that verifies the request size is within limits.
@@ -403,8 +405,7 @@ class Pathod(tcp.TCPServer):
 
     def add_log(self, d):
         if not self.noapi:
-            lock = threading.Lock()
-            with lock:
+            with self.loglock:
                 d["id"] = self.logid
                 self.log.insert(0, d)
                 if len(self.log) > self.LOGBUF:
@@ -413,17 +414,18 @@ class Pathod(tcp.TCPServer):
             return d["id"]
 
     def clear_log(self):
-        lock = threading.Lock()
-        with lock:
+        with self.loglock:
             self.log = []
 
     def log_by_id(self, identifier):
-        for i in self.log:
-            if i["id"] == identifier:
-                return i
+        with self.loglock:
+            for i in self.log:
+                if i["id"] == identifier:
+                    return i
 
     def get_log(self):
-        return self.log
+        with self.loglock:
+            return self.log
 
 
 def main(args):  # pragma: no cover
diff --git a/pathod/protocols/websockets.py b/pathod/protocols/websockets.py
index 134d27bc..2b60e618 100644
--- a/pathod/protocols/websockets.py
+++ b/pathod/protocols/websockets.py
@@ -18,7 +18,7 @@ class WebsocketsProtocol:
                     frm = websockets.Frame.from_file(self.pathod_handler.rfile)
                 except NetlibException as e:
                     lg("Error reading websocket frame: %s" % e)
-                    break
+                    return None, None
                 ended = time.time()
                 lg(frm.human_readable())
             retlog = dict(
diff --git a/pathod/test.py b/pathod/test.py
index 23b7a5b6..11462729 100644
--- a/pathod/test.py
+++ b/pathod/test.py
@@ -1,12 +1,14 @@
 from six.moves import cStringIO as StringIO
 import threading
+import time
+
 from six.moves import queue
 
-import requests
-import requests.packages.urllib3
 from . import pathod
 
-requests.packages.urllib3.disable_warnings()
+
+class TimeoutError(Exception):
+    pass
 
 
 class Daemon:
@@ -39,39 +41,51 @@ class Daemon:
         """
         return "%s/p/%s" % (self.urlbase, spec)
 
-    def info(self):
-        """
-            Return some basic info about the remote daemon.
-        """
-        resp = requests.get("%s/api/info" % self.urlbase, verify=False)
-        return resp.json()
-
     def text_log(self):
         return self.logfp.getvalue()
 
+    def wait_for_silence(self, timeout=5):
+        start = time.time()
+        while 1:
+            if time.time() - start >= timeout:
+                raise TimeoutError(
+                    "%s service threads still alive" %
+                    self.thread.server.handler_counter.count
+                )
+            if self.thread.server.handler_counter.count == 0:
+                return
+
+    def expect_log(self, n, timeout=5):
+        l = []
+        start = time.time()
+        while True:
+            l = self.log()
+            if time.time() - start >= timeout:
+                return None
+            if len(l) >= n:
+                break
+        return l
+
     def last_log(self):
         """
             Returns the last logged request, or None.
         """
-        l = self.log()
+        l = self.expect_log(1)
         if not l:
             return None
-        return l[0]
+        return l[-1]
 
     def log(self):
         """
             Return the log buffer as a list of dictionaries.
         """
-        resp = requests.get("%s/api/log" % self.urlbase, verify=False)
-        return resp.json()["log"]
+        return self.thread.server.get_log()
 
     def clear_log(self):
         """
             Clear the log.
         """
-        self.logfp.truncate(0)
-        resp = requests.get("%s/api/clear_log" % self.urlbase, verify=False)
-        return resp.ok
+        return self.thread.server.clear_log()
 
     def shutdown(self):
         """
@@ -88,6 +102,7 @@ class _PaThread(threading.Thread):
         self.name = "PathodThread"
         self.iface, self.q, self.ssl = iface, q, ssl
         self.daemonargs = daemonargs
+        self.server = None
 
     def run(self):
         self.server = pathod.Pathod(
diff --git a/setup.cfg b/setup.cfg
index e83ca0ab..eeaac0c8 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -11,8 +11,7 @@ addopts = --capture=no
 
 [coverage:run]
 branch = True
-include = *mitmproxy*, *netlib*, *pathod*
-omit = *contrib*, *tnetstring*, *platform*, *console*, *main.py
+omit = *contrib*, *tnetstring*, *platform*, *main.py
 
 [coverage:report]
 show_missing = True
diff --git a/test/pathod/test_app.py b/test/pathod/test_app.py
index ac89c44c..fbaa773c 100644
--- a/test/pathod/test_app.py
+++ b/test/pathod/test_app.py
@@ -11,11 +11,11 @@ class TestApp(tutils.DaemonTests):
 
     def test_about(self):
         r = self.getpath("/about")
-        assert r.ok
+        assert r.status_code == 200
 
     def test_download(self):
         r = self.getpath("/download")
-        assert r.ok
+        assert r.status_code == 200
 
     def test_docs(self):
         assert self.getpath("/docs/pathod").status_code == 200
@@ -27,7 +27,7 @@ class TestApp(tutils.DaemonTests):
     def test_log(self):
         assert self.getpath("/log").status_code == 200
         assert self.get("200:da").status_code == 200
-        id = self.d.log()[0]["id"]
+        id = self.d.expect_log(1)[0]["id"]
         assert self.getpath("/log").status_code == 200
         assert self.getpath("/log/%s" % id).status_code == 200
         assert self.getpath("/log/9999999").status_code == 404
diff --git a/test/pathod/test_pathod.py b/test/pathod/test_pathod.py
index 4d969158..ec9c169f 100644
--- a/test/pathod/test_pathod.py
+++ b/test/pathod/test_pathod.py
@@ -1,7 +1,6 @@
 from six.moves import cStringIO as StringIO
-import pytest
 
-from pathod import pathod, version
+from pathod import pathod
 from netlib import tcp
 from netlib.exceptions import HttpException, TlsException
 import tutils
@@ -129,7 +128,6 @@ class CommonTests(tutils.DaemonTests):
         assert self.d.last_log()
         # FIXME: Other binary data elements
 
-    @pytest.mark.skip(reason="race condition")
     def test_sizelimit(self):
         r = self.get("200:b@1g")
         assert r.status_code == 800
@@ -140,21 +138,15 @@ class CommonTests(tutils.DaemonTests):
         r, _ = self.pathoc([r"get:'/p/200':i0,'\r\n'"])
         assert r[0].status_code == 200
 
-    def test_info(self):
-        assert tuple(self.d.info()["version"]) == version.IVERSION
-
-    @pytest.mark.skip(reason="race condition")
     def test_logs(self):
-        assert self.d.clear_log()
-        assert not self.d.last_log()
+        self.d.clear_log()
         assert self.get("202:da")
-        assert len(self.d.log()) == 1
-        assert self.d.clear_log()
+        assert self.d.expect_log(1)
+        self.d.clear_log()
         assert len(self.d.log()) == 0
 
     def test_disconnect(self):
-        rsp = self.get("202:b@100k:d200")
-        assert len(rsp.content) < 200
+        tutils.raises("unexpected eof", self.get, "202:b@100k:d200")
 
     def test_parserr(self):
         rsp = self.get("400:msg,b:")
@@ -166,7 +158,7 @@ class CommonTests(tutils.DaemonTests):
         assert rsp.content.strip() == "testfile"
 
     def test_anchor(self):
-        rsp = self.getpath("anchor/foo")
+        rsp = self.getpath("/anchor/foo")
         assert rsp.status_code == 202
 
     def test_invalid_first_line(self):
@@ -223,7 +215,6 @@ class CommonTests(tutils.DaemonTests):
         )
         assert r[1].payload == "test"
 
-    @pytest.mark.skip(reason="race condition")
     def test_websocket_frame_reflect_error(self):
         r, _ = self.pathoc(
             ["ws:/p/", "wf:-mask:knone:f'wf:b@10':i13,'a'"],
@@ -233,7 +224,6 @@ class CommonTests(tutils.DaemonTests):
         # FIXME: Race Condition?
         assert "Parse error" in self.d.text_log()
 
-    @pytest.mark.skip(reason="race condition")
     def test_websocket_frame_disconnect_error(self):
         self.pathoc(["ws:/p/", "wf:b@10:d3"], ws_read_limit=0)
         assert self.d.last_log()
diff --git a/test/pathod/test_test.py b/test/pathod/test_test.py
index cee286a4..6399894e 100644
--- a/test/pathod/test_test.py
+++ b/test/pathod/test_test.py
@@ -2,6 +2,10 @@ import logging
 import requests
 from pathod import test
 import tutils
+
+import requests.packages.urllib3
+
+requests.packages.urllib3.disable_warnings()
 logging.disable(logging.CRITICAL)
 
 
diff --git a/test/pathod/tutils.py b/test/pathod/tutils.py
index f7bb22e5..b9f38d86 100644
--- a/test/pathod/tutils.py
+++ b/test/pathod/tutils.py
@@ -3,6 +3,7 @@ import re
 import shutil
 import requests
 from six.moves import cStringIO as StringIO
+import urllib
 
 from netlib import tcp
 from netlib import utils
@@ -63,10 +64,11 @@ class DaemonTests(object):
         shutil.rmtree(cls.confdir)
 
     def teardown(self):
+        self.d.wait_for_silence()
         if not (self.noweb or self.noapi):
             self.d.clear_log()
 
-    def getpath(self, path, params=None):
+    def _getpath(self, path, params=None):
         scheme = "https" if self.ssl else "http"
         resp = requests.get(
             "%s://localhost:%s/%s" % (
@@ -79,9 +81,29 @@ class DaemonTests(object):
         )
         return resp
 
+    def getpath(self, path, params=None):
+        logfp = StringIO()
+        c = pathoc.Pathoc(
+            ("localhost", self.d.port),
+            ssl=self.ssl,
+            fp=logfp,
+        )
+        with c.connect():
+            if params:
+                path = path + "?" + urllib.urlencode(params)
+            resp = c.request("get:%s" % path)
+            return resp
+
     def get(self, spec):
-        resp = requests.get(self.d.p(spec), verify=False)
-        return resp
+        logfp = StringIO()
+        c = pathoc.Pathoc(
+            ("localhost", self.d.port),
+            ssl=self.ssl,
+            fp=logfp,
+        )
+        with c.connect():
+            resp = c.request("get:/p/%s" % urllib.quote(spec).encode("string_escape"))
+            return resp
 
     def pathoc(
         self,
@@ -106,16 +128,16 @@ class DaemonTests(object):
             fp=logfp,
             use_http2=use_http2,
         )
-        c.connect(connect_to)
-        ret = []
-        for i in specs:
-            resp = c.request(i)
-            if resp:
-                ret.append(resp)
-        for frm in c.wait():
-            ret.append(frm)
-        c.stop()
-        return ret, logfp.getvalue()
+        with c.connect(connect_to):
+            ret = []
+            for i in specs:
+                resp = c.request(i)
+                if resp:
+                    ret.append(resp)
+            for frm in c.wait():
+                ret.append(frm)
+            c.stop()
+            return ret, logfp.getvalue()
 
 
 tmpdir = tutils.tmpdir