Skip to content

Commit

Permalink
Black formmatted
Browse files Browse the repository at this point in the history
  • Loading branch information
wjdecorte committed Aug 13, 2019
1 parent 55efe33 commit 7580b23
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 81 deletions.
8 changes: 6 additions & 2 deletions pydruid/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,15 @@
import StringIO
from gzip import GzipFile


def decompress(data):
infile = StringIO.StringIO()
infile.write(data)
with GzipFile(fileobj=infile, mode="r") as f:
f.rewind()
ud = f.read()
return ud


else:
from gzip import compress, decompress

Expand Down Expand Up @@ -560,6 +561,7 @@ class PyDruid(BaseDruidClient):
0 7 2013-10-04T00:00:00.000Z user_1
1 6 2013-10-04T00:00:00.000Z user_2
"""

def __init__(self, url, endpoint, extra_headers=None):
super(PyDruid, self).__init__(url, endpoint, extra_headers)

Expand All @@ -572,7 +574,9 @@ def _post(self, query):
if content_encoding == "gzip":
data = decompress(res.read()).decode("utf-8")
elif content_encoding:
raise ValueError("Invalid content encoding: {}".format(content_encoding))
raise ValueError(
"Invalid content encoding: {}".format(content_encoding)
)
else:
data = res.read().decode("utf-8")
res.close()
Expand Down
175 changes: 96 additions & 79 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,13 @@
if sys.version_info.major == 2 and sys.version_info.minor == 7:
from gzip import GzipFile


def compress(data):
out = StringIO()
with GzipFile(fileobj=out, mode="w") as f:
f.write(data)
return out.getvalue()


else:
from gzip import compress, decompress

Expand All @@ -30,23 +31,19 @@ def create_client(headers=None):


def create_blank_query():
return Query({}, 'none')
return Query({}, "none")


def _http_error(code, msg, data = ''):
def _http_error(code, msg, data=""):
# Need a file-like object for the response data
fp = StringIO(data)
return urllib.error.HTTPError(
url='http://fakeurl:8080/druid/v2/',
hdrs={},
code=code,
msg=msg,
fp=fp,
url="http://fakeurl:8080/druid/v2/", hdrs={}, code=code, msg=msg, fp=fp
)


class TestPyDruid:
@patch('pydruid.client.urllib.request.urlopen')
@patch("pydruid.client.urllib.request.urlopen")
def test_druid_returns_error(self, mock_urlopen):
# given
mock_urlopen.side_effect = _http_error(500, "Druid error")
Expand All @@ -55,20 +52,22 @@ def test_druid_returns_error(self, mock_urlopen):
# when / then
with pytest.raises(IOError):
client.topn(
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000})

@patch('pydruid.client.urllib.request.urlopen')
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000},
)

@patch("pydruid.client.urllib.request.urlopen")
def test_druid_returns_html_error(self, mock_urlopen):
# given
message = textwrap.dedent("""
message = textwrap.dedent(
"""
<html>
<head>
<meta http-equiv="Content-Type" content="text/html;charset=ISO-8859-1"/>
Expand All @@ -81,26 +80,31 @@ def test_druid_returns_html_error(self, mock_urlopen):
<hr /><a href="http://eclipse.org/jetty">Powered by Jetty:// 9.3.19.v20170502</a><hr/>
</body>
</html>
""").strip()
mock_urlopen.side_effect = _http_error(500, 'Internal Server Error', message)
"""
).strip()
mock_urlopen.side_effect = _http_error(500, "Internal Server Error", message)
client = create_client()

# when / then
with pytest.raises(IOError) as e:
client.topn(
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000})

assert str(e.value) == textwrap.dedent("""
HTTP Error 500: Internal Server Error
Druid Error: javax.servlet.ServletException: java.lang.OutOfMemoryError: GC overhead limit exceeded
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000},
)

assert (
str(e.value)
== textwrap.dedent(
"""
HTTP Error 500: Internal Server Error
Druid Error: javax.servlet.ServletException: java.lang.OutOfMemoryError: GC overhead limit exceeded
Query is: {
"aggregations": [
{
Expand All @@ -125,9 +129,11 @@ def test_druid_returns_html_error(self, mock_urlopen):
"queryType": "topN",
"threshold": 1
}
""").strip()
"""
).strip()
)

@patch('pydruid.client.urllib.request.urlopen')
@patch("pydruid.client.urllib.request.urlopen")
def test_druid_returns_results(self, mock_urlopen):
# given
response = Mock()
Expand All @@ -139,29 +145,32 @@ def test_druid_returns_results(self, mock_urlopen):
"metric" : 100
} ]
} ]
""".encode("utf-8")
""".encode(
"utf-8"
)
response.info.return_value = {}
mock_urlopen.return_value = response
client = create_client()

# when
top = client.topn(
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000})
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000},
)

# then
assert top is not None
assert len(top.result) == 1
assert len(top.result[0]['result']) == 1
assert len(top.result[0]["result"]) == 1

@patch('pydruid.client.urllib.request.urlopen')
@patch("pydruid.client.urllib.request.urlopen")
def test_client_allows_to_export_last_query(self, mock_urlopen):
# given
response = Mock()
Expand All @@ -173,70 +182,78 @@ def test_client_allows_to_export_last_query(self, mock_urlopen):
"metric" : 100
} ]
} ]
""".encode("utf-8")
""".encode(
"utf-8"
)
response.info.return_value = {}
mock_urlopen.return_value = response
client = create_client()
client.topn(
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000})
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000},
)

# when / then
# assert that last_query.export_tsv method was called (it should throw an exception, given empty path)
with pytest.raises(TypeError):
client.export_tsv(None)

@patch('pydruid.client.urllib.request.urlopen')
@patch("pydruid.client.urllib.request.urlopen")
def test_client_auth_creds(self, mock_urlopen):
client = create_client()
query = create_blank_query()
client.set_basic_auth_credentials('myUsername', 'myPassword')
client.set_basic_auth_credentials("myUsername", "myPassword")
headers, _, _ = client._prepare_url_headers_and_body(query)
assert headers['Authorization'] == "Basic bXlVc2VybmFtZTpteVBhc3N3b3Jk"
assert headers["Authorization"] == "Basic bXlVc2VybmFtZTpteVBhc3N3b3Jk"

def test_client_allows_extra_headers(self):
client = create_client(headers={'Accept-Encoding': 'gzip'})
client = create_client(headers={"Accept-Encoding": "gzip"})
query = create_blank_query()
headers, _, _ = client._prepare_url_headers_and_body(query)
assert headers['Accept-Encoding'] == "gzip"
assert headers["Accept-Encoding"] == "gzip"

@patch('pydruid.client.urllib.request.urlopen')
@patch("pydruid.client.urllib.request.urlopen")
def test_return_compressed_data(self, mock_urlopen):
# given
response = Mock()
response.read.return_value = compress("""
response.read.return_value = compress(
"""
[ {
"timestamp" : "2015-12-30T14:14:49.000Z",
"result" : [ {
"dimension" : "aaaa",
"metric" : 100
} ]
} ]
""".encode("utf-8"))
""".encode(
"utf-8"
)
)
response.info.return_value = {"Content-Encoding": "gzip"}
mock_urlopen.return_value = response
client = create_client(headers={'Accept-Encoding': 'gzip'})
client = create_client(headers={"Accept-Encoding": "gzip"})

# when
top = client.topn(
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000})
datasource="testdatasource",
granularity="all",
intervals="2015-12-29/pt1h",
aggregations={"count": doublesum("count")},
dimension="user_name",
metric="count",
filter=Dimension("user_lang") == "en",
threshold=1,
context={"timeout": 1000},
)

# then
assert top is not None
assert len(top.result) == 1
assert len(top.result[0]['result']) == 1
assert len(top.result[0]["result"]) == 1

0 comments on commit 7580b23

Please sign in to comment.