diff --git a/mitmproxy2swagger/test_mitmproxy2swagger.py b/mitmproxy2swagger/test_mitmproxy2swagger.py
index 5c8a652..4058d09 100644
--- a/mitmproxy2swagger/test_mitmproxy2swagger.py
+++ b/mitmproxy2swagger/test_mitmproxy2swagger.py
@@ -1,77 +1,9 @@
 # -*- coding: utf-8 -*-
 import sys
-import tempfile
-from typing import Any, List
 
 import ruamel.yaml as ruamel
 
-from .mitmproxy2swagger import main
-
-
-def get_nested_key(obj: Any, path: str) -> Any:
-    """Gets a nested key from a dict."""
-    keys = path.split(".")
-    for key in keys:
-        if not isinstance(obj, dict):
-            return None
-        if key not in obj:
-            return None
-        obj = obj[key]
-    return obj
-
-
-def mitmproxy2swagger_e2e_test(
-    input_file: str, url_prefix: str, extra_args: List[str] | None = None
-) -> Any:
-    """Runs mitmproxy2swagger on the given input file twice, and returns the
-    detected endpoints."""
-    yaml_tmp_path = tempfile.mktemp(suffix=".yaml", prefix="sklep.lisek.")
-    main(
-        [
-            "-i",
-            input_file,
-            "-o",
-            yaml_tmp_path,
-            "-p",
-            url_prefix,
-        ]
-        + (extra_args or [])
-    )
-    yaml = ruamel.YAML()
-
-    data = None
-    # try to parse the file
-    with open(yaml_tmp_path, "r") as f:
-        data = yaml.load(f.read())
-    assert data is not None
-    assert "x-path-templates" in data
-    assert "servers" in data
-    # remove all of the ignore:prefixes in x-path-templates
-    data["x-path-templates"] = [
-        x.replace("ignore:", "") for x in data["x-path-templates"]
-    ]
-
-    # save the file
-    with open(yaml_tmp_path, "w") as f:
-        yaml.dump(data, f)
-
-    # run mitmproxy2swagger again
-    main(
-        [
-            "-i",
-            input_file,
-            "-o",
-            yaml_tmp_path,
-            "-p",
-            url_prefix,
-        ]
-        + (extra_args or [])
-    )
-
-    # load the file again
-    with open(yaml_tmp_path, "r") as f:
-        data = yaml.load(f.read())
-    return data
+from mitmproxy2swagger.test_util import get_nested_key, mitmproxy2swagger_e2e_test
 
 
 def test_mitmproxy2swagger_generates_swagger_from_har():
@@ -136,4 +68,7 @@ def test_mitmproxy2swagger_generates_headers_for_flow_files():
         ],
     )
     assert data is not None
-    assert get_nested_key(data, "paths./post.post.responses.200.headers.content-type") is not None
+    assert (
+        get_nested_key(data, "paths./post.post.responses.200.headers.content-type")
+        is not None
+    )
diff --git a/mitmproxy2swagger/test_util.py b/mitmproxy2swagger/test_util.py
new file mode 100644
index 0000000..297fc32
--- /dev/null
+++ b/mitmproxy2swagger/test_util.py
@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+
+import tempfile
+from typing import Any, List
+
+import ruamel.yaml as ruamel
+
+from .mitmproxy2swagger import main
+
+
+def get_nested_key(obj: Any, path: str) -> Any:
+    """Gets a nested key from a dict."""
+    keys = path.split(".")
+    for key in keys:
+        if not isinstance(obj, dict):
+            return None
+        if key not in obj:
+            return None
+        obj = obj[key]
+    return obj
+
+
+def mitmproxy2swagger_e2e_test(
+    input_file: str, url_prefix: str, extra_args: List[str] | None = None
+) -> Any:
+    """Runs mitmproxy2swagger on the given input file twice, and returns the
+    detected endpoints."""
+    yaml_tmp_path = tempfile.mktemp(suffix=".yaml", prefix="sklep.lisek.")
+    main(
+        [
+            "-i",
+            input_file,
+            "-o",
+            yaml_tmp_path,
+            "-p",
+            url_prefix,
+        ]
+        + (extra_args or [])
+    )
+    yaml = ruamel.YAML()
+
+    data = None
+    # try to parse the file
+    with open(yaml_tmp_path, "r") as f:
+        data = yaml.load(f.read())
+    assert data is not None
+    assert "x-path-templates" in data
+    assert "servers" in data
+    # remove all of the ignore:prefixes in x-path-templates
+    data["x-path-templates"] = [
+        x.replace("ignore:", "") for x in data["x-path-templates"]
+    ]
+
+    # save the file
+    with open(yaml_tmp_path, "w") as f:
+        yaml.dump(data, f)
+
+    # run mitmproxy2swagger again
+    main(
+        [
+            "-i",
+            input_file,
+            "-o",
+            yaml_tmp_path,
+            "-p",
+            url_prefix,
+        ]
+        + (extra_args or [])
+    )
+
+    # load the file again
+    with open(yaml_tmp_path, "r") as f:
+        data = yaml.load(f.read())
+    return data