diff --git a/gvm/transforms.py b/gvm/transforms.py index 7add892da..5f4b42edf 100644 --- a/gvm/transforms.py +++ b/gvm/transforms.py @@ -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 @@ -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) @@ -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) @@ -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) diff --git a/tests/transforms/test_check_command_transform.py b/tests/transforms/test_check_command_transform.py index 15615a8e5..0ea4ddb4c 100644 --- a/tests/transforms/test_check_command_transform.py +++ b/tests/transforms/test_check_command_transform.py @@ -16,7 +16,7 @@ def test_no_status_transform(self): transform = CheckCommandTransform() with self.assertRaises(GvmError): - transform("") + transform(b"") def test_no_success_300status_transform(self): transform = CheckCommandTransform() @@ -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) @@ -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) @@ -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) @@ -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), '') + self.assertEqual(transform(response), b'') diff --git a/tests/transforms/test_etree_check_command_transform.py b/tests/transforms/test_etree_check_command_transform.py index 61e04e704..38e0b0deb 100644 --- a/tests/transforms/test_etree_check_command_transform.py +++ b/tests/transforms/test_etree_check_command_transform.py @@ -16,7 +16,7 @@ def test_no_status_transform(self): transform = EtreeCheckCommandTransform() with self.assertRaises(GvmError): - transform("") + transform(b"") def test_no_success_status_transform(self): transform = EtreeCheckCommandTransform() @@ -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) @@ -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) diff --git a/tests/transforms/test_etree_transform.py b/tests/transforms/test_etree_transform.py index feeb11e56..b4dfe08e0 100644 --- a/tests/transforms/test_etree_transform.py +++ b/tests/transforms/test_etree_transform.py @@ -13,13 +13,13 @@ class EtreeTransformTestCase(unittest.TestCase): def test_transform_response(self): transform = EtreeTransform() - result = transform("") + result = transform(b"") self.assertTrue(etree.iselement(result)) def test_transform_more_complex_response(self): transform = EtreeTransform() - result = transform('') + result = transform(b'') self.assertTrue(etree.iselement(result)) self.assertEqual(result.tag, "foo")