Ruff cleanup

2 files changed, 11 insertions, 9 deletions

diff --git a/tests/models/network/test_nmap.py b/tests/models/network/test_nmap.py
index 0be98d4..a135a13 100644
--- a/tests/models/network/test_nmap.py
+++ b/tests/models/network/test_nmap.py
@@ -1,20 +1,21 @@
 import subprocess
 
-from generalresearch.models.network.definitions import IPProtocol
-from generalresearch.models.network.nmap.execute import execute_nmap
 import faker
 
+from generalresearch.managers.network.tool_run import ToolRunManager
+from generalresearch.models.network.definitions import IPProtocol
+from generalresearch.models.network.nmap.execute import execute_nmap
 from generalresearch.models.network.nmap.result import PortState
-from generalresearch.models.network.tool_run import ToolName, ToolClass
+from generalresearch.models.network.tool_run import ToolClass, ToolName
 
 fake = faker.Faker()
 
 
-def resolve(host):
+def resolve(host: str):
     return subprocess.check_output(["dig", host, "+short"]).decode().strip()
 
 
-def test_execute_nmap_scanme(toolrun_manager):
+def test_execute_nmap_scanme(toolrun_manager: ToolRunManager):
     ip = resolve("scanme.nmap.org")
 
     run = execute_nmap(ip=ip, top_ports=None, ports="20-30", enable_advanced=False)
diff --git a/tests/models/network/test_rdns.py b/tests/models/network/test_rdns.py
index e56c494..5c3b024 100644
--- a/tests/models/network/test_rdns.py
+++ b/tests/models/network/test_rdns.py
@@ -1,12 +1,13 @@
-from generalresearch.models.network.rdns.execute import execute_rdns
 import faker
 
-from generalresearch.models.network.tool_run import ToolName, ToolClass
+from generalresearch.managers.network.tool_run import ToolRunManager
+from generalresearch.models.network.rdns.execute import execute_rdns
+from generalresearch.models.network.tool_run import ToolClass, ToolName
 
 fake = faker.Faker()
 
 
-def test_execute_rdns_grl(toolrun_manager):
+def test_execute_rdns_grl(toolrun_manager: ToolRunManager):
     ip = "65.19.129.53"
     run = execute_rdns(ip=ip)
     assert run.tool_name == ToolName.DIG
@@ -20,7 +21,7 @@ def test_execute_rdns_grl(toolrun_manager):
     toolrun_manager.create_rdns_run(run)
 
 
-def test_execute_rdns_none(toolrun_manager):
+def test_execute_rdns_none(toolrun_manager: ToolRunManager):
     ip = fake.ipv6()
     run = execute_rdns(ip)
     result = run.parsed