Commit ec4f2138 authored by Alexandre's avatar Alexandre
Browse files

Check for expected string as a class method

parent 73ccdb9c
......@@ -31,6 +31,12 @@ class Request:
def to_wire(self):
self.data = self.message.to_wire()
def has_expected_str(self, string):
try:
return string is None or string in str(self.response)
except AttributeError:
return False
class RequestDOT(Request):
# raising custom exception for each unexpected response might be a good idea
......
......@@ -89,7 +89,7 @@ def print_result(connection, request, prefix=None, display_err=True):
msg = request.response
size = request.response_size
if (dot and rcode) or (not dot and rcode == 200):
if opts.expect is not None and opts.expect not in str(request.response):
if not request.has_expected_str(opts.expect):
print("%s Cannot find \"%s\" in response" % (server, opts.expect))
ok = False
if ok and size is not None and size > 0:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment