Skip to content

Commit

Permalink
Adjust transform classes for bytes based protocol
Browse files Browse the repository at this point in the history
For the protocol implementation bytes are passed in the form of
Responses to the transform callables. Therefore adjust the Transform
classes to handle bytes now.
  • Loading branch information
bjoernricks committed Jun 18, 2024
1 parent 2bc27a2 commit f2d7801
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 16 deletions.
12 changes: 7 additions & 5 deletions gvm/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
"""
from lxml import etree

from gvm.protocols.core import Response

from .errors import GvmError, GvmResponseError, GvmServerError
from .xml import Element, create_parser

Expand All @@ -19,10 +21,10 @@ class EtreeTransform:
def __init__(self):
self._parser = create_parser()

def _convert_response(self, response: str) -> Element:
return etree.XML(response, parser=self._parser)
def _convert_response(self, response: Response) -> Element:
return etree.XML(bytes(response), parser=self._parser)

def __call__(self, response: str) -> Element:
def __call__(self, response: Response) -> Element:
return self._convert_response(response)


Expand All @@ -46,7 +48,7 @@ class CheckCommandTransform(EtreeTransform):
response was an error response
"""

def __call__(self, response: str) -> str: # type: ignore[override]
def __call__(self, response: Response) -> Response: # type: ignore[override]
root = self._convert_response(response)

check_command_status(root)
Expand All @@ -60,7 +62,7 @@ class EtreeCheckCommandTransform(EtreeTransform):
response was an error response
"""

def __call__(self, response: str) -> Element:
def __call__(self, response: Response) -> Element:
root = self._convert_response(response)

check_command_status(root)
Expand Down
12 changes: 6 additions & 6 deletions tests/transforms/test_check_command_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def test_no_status_transform(self):
transform = CheckCommandTransform()

with self.assertRaises(GvmError):
transform("<foo/>")
transform(b"<foo/>")

def test_no_success_300status_transform(self):
transform = CheckCommandTransform()
Expand All @@ -25,7 +25,7 @@ def test_no_success_300status_transform(self):
root.set("status", "300")
root.set("status_text", "Foo error")

response = etree.tostring(root).decode("utf-8")
response = etree.tostring(root)

with self.assertRaises(GvmError):
transform(response)
Expand All @@ -37,7 +37,7 @@ def test_no_success_400status_transform(self):
root.set("status", "400")
root.set("status_text", "Foo error")

response = etree.tostring(root).decode("utf-8")
response = etree.tostring(root)

with self.assertRaises(GvmResponseError):
transform(response)
Expand All @@ -49,7 +49,7 @@ def test_no_success_500status_transform(self):
root.set("status", "500")
root.set("status_text", "Foo error")

response = etree.tostring(root).decode("utf-8")
response = etree.tostring(root)

with self.assertRaises(GvmServerError):
transform(response)
Expand All @@ -60,6 +60,6 @@ def test_success_response(self):
root = etree.Element("foo_response")
root.set("status", "200")

response = etree.tostring(root).decode("utf-8")
response = etree.tostring(root)

self.assertEqual(transform(response), '<foo_response status="200"/>')
self.assertEqual(transform(response), b'<foo_response status="200"/>')
6 changes: 3 additions & 3 deletions tests/transforms/test_etree_check_command_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def test_no_status_transform(self):
transform = EtreeCheckCommandTransform()

with self.assertRaises(GvmError):
transform("<foo/>")
transform(b"<foo/>")

def test_no_success_status_transform(self):
transform = EtreeCheckCommandTransform()
Expand All @@ -25,7 +25,7 @@ def test_no_success_status_transform(self):
root.set("status", "400")
root.set("status_text", "Foo error")

response = etree.tostring(root).decode("utf-8")
response = etree.tostring(root)

with self.assertRaises(GvmError):
transform(response)
Expand All @@ -36,7 +36,7 @@ def test_success_response(self):
root = etree.Element("foo_response")
root.set("status", "200")

response = etree.tostring(root).decode("utf-8")
response = etree.tostring(root)

result = transform(response)

Expand Down
4 changes: 2 additions & 2 deletions tests/transforms/test_etree_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,13 @@
class EtreeTransformTestCase(unittest.TestCase):
def test_transform_response(self):
transform = EtreeTransform()
result = transform("<foo/>")
result = transform(b"<foo/>")

self.assertTrue(etree.iselement(result))

def test_transform_more_complex_response(self):
transform = EtreeTransform()
result = transform('<foo id="bar"><lorem/><ipsum/></foo>')
result = transform(b'<foo id="bar"><lorem/><ipsum/></foo>')

self.assertTrue(etree.iselement(result))
self.assertEqual(result.tag, "foo")
Expand Down

0 comments on commit f2d7801

Please sign in to comment.