Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adjust transform classes for bytes based protocol #1151

Merged
merged 1 commit into from
Jun 18, 2024
Merged

Adjust transform classes for bytes based protocol #1151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions gvm/transforms.py
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
Loading