Ruff afternoon!

4 files changed, 24 insertions, 6 deletions

diff --git a/tests/models/network/test_mtr.py b/tests/models/network/test_mtr.py
index 840a773..7f8a736 100644
--- a/tests/models/network/test_mtr.py
+++ b/tests/models/network/test_mtr.py
@@ -1,12 +1,15 @@
+from __future__ import annotations
+
 import faker
 
+from generalresearch.managers.network.tool_run import ToolRunManager
 from generalresearch.models.network.mtr.execute import execute_mtr
 from generalresearch.models.network.tool_run import ToolClass, ToolName
 
 fake = faker.Faker()
 
 
-def test_execute_mtr(toolrun_manager):
+def test_execute_mtr(toolrun_manager: ToolRunManager):
     ip = "65.19.129.53"
 
     run = execute_mtr(ip=ip, report_cycles=3)
diff --git a/tests/models/network/test_nmap.py b/tests/models/network/test_nmap.py
index a135a13..5e9f4d0 100644
--- a/tests/models/network/test_nmap.py
+++ b/tests/models/network/test_nmap.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import subprocess
 
 import faker
@@ -5,8 +7,8 @@ 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 ToolClass, ToolName
+from generalresearch.models.network.nmap.result import NmapResult, PortState
+from generalresearch.models.network.tool_run import NmapRun, Status, ToolClass, ToolName
 
 fake = faker.Faker()
 
@@ -18,10 +20,13 @@ def resolve(host: str):
 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)
+    run: NmapRun = execute_nmap(
+        ip=ip, top_ports=None, ports="20-30", enable_advanced=False
+    )
     assert run.tool_name == ToolName.NMAP
     assert run.tool_class == ToolClass.PORT_SCAN
     assert run.ip == ip
+    assert isinstance(run.parsed, NmapResult)
     result = run.parsed
 
     port22 = result._port_index[(IPProtocol.TCP, 22)]
diff --git a/tests/models/network/test_nmap_parser.py b/tests/models/network/test_nmap_parser.py
index 7822380..473a63f 100644
--- a/tests/models/network/test_nmap_parser.py
+++ b/tests/models/network/test_nmap_parser.py
@@ -1,8 +1,14 @@
+from __future__ import annotations
+
 import os
 
 import pytest
 
 from generalresearch.models.network.nmap.parser import parse_nmap_xml
+from generalresearch.models.network.nmap.result import (
+    NmapResult,
+    NmapTrace,
+)
 
 
 @pytest.fixture
@@ -13,9 +19,11 @@ def nmap_raw_output_2(request) -> str:
     return data
 
 
-def test_nmap_xml_parser(nmap_raw_output, nmap_raw_output_2):
-    n = parse_nmap_xml(nmap_raw_output)
+def test_nmap_xml_parser(nmap_raw_output: str, nmap_raw_output_2: str):
+    n: NmapResult = parse_nmap_xml(nmap_raw_output)
     assert n.tcp_open_ports == [61232]
+
+    assert isinstance(n.trace, NmapTrace)
     assert len(n.trace.hops) == 18
 
     n = parse_nmap_xml(nmap_raw_output_2)
diff --git a/tests/models/network/test_rdns.py b/tests/models/network/test_rdns.py
index 5c3b024..1a15a28 100644
--- a/tests/models/network/test_rdns.py
+++ b/tests/models/network/test_rdns.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import faker
 
 from generalresearch.managers.network.tool_run import ToolRunManager