diff --git a/pydu/request.py b/pydu/request.py index b4518db..834ed0d 100644 --- a/pydu/request.py +++ b/pydu/request.py @@ -1,6 +1,9 @@ import os import shutil import tempfile +import socket + +from pydu import logger from pydu.compat import PY2, string_types, urlparse, ulib from pydu.string import safeunicode @@ -97,3 +100,24 @@ def download(url, dst=None): shutil.move(tmpfile, filename) return filename + + +def check_connect(ip, port, retry=1, timeout=0.5): + while retry: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + except socket.error as e: + logger.exception(e) + retry -= 1 + continue + + try: + s.settimeout(timeout) + s.connect((ip, port)) + return s.getsockname()[0] + except socket.error: + logger.error("Connect to ip:%s port:%d fail", ip, port) + s.close() + finally: + retry -= 1 + return None diff --git a/requirements-dev.txt b/requirements-dev.txt index 60551a5..1f81207 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,2 +1,3 @@ pytest -coveralls \ No newline at end of file +coveralls +HTTPretty==0.8.14 \ No newline at end of file diff --git a/tests/test_request.py b/tests/test_request.py index 3e29a62..183cb30 100644 --- a/tests/test_request.py +++ b/tests/test_request.py @@ -1,4 +1,5 @@ -from pydu.request import FileName +import httpretty +from pydu.request import FileName, check_connect def test_filename_from_url(): @@ -36,3 +37,10 @@ def test_filename_from_headers(): headers = 'Content-Disposition: attachment; filename=' assert FileName.from_headers(headers) is None + + +@httpretty.activate +def test_check_connect(): + httpretty.register_uri(httpretty.GET, 'http://localhost') + assert check_connect('http://localhost', port=80, timeout=0.01) + assert not check_connect('http://localhost', port=8000, timeout=0.01)