tests: add zmap tests

.github/workflows/ci.yml

@@ -34,4 +34,8 @@ jobs:
 
     - name: Run linting
       run: |
-        make lint
+        make lint
+
+    - name: Run tests
+      run: |
+        make pytest

Makefile

@@ -19,3 +19,5 @@ clean:
 	rm -rf *.egg-info
 	find . -type d -name __pycache__ -exec rm -rf {} +
 	find . -type f -name "*.pyc" -delete
+pytest:
+	poetry run pytest

tests/test_api.py

@@ -0,0 +1,70 @@
+from collections.abc import Generator
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+from starlette.testclient import TestClient
+
+from zmapsdk.api import app
+
+
+@pytest.fixture
+def client() -> Generator[TestClient, Any, None]:
+    """Fixture that provides a test client for FastAPI"""
+    with TestClient(app) as test_client:
+        yield test_client
+
+
+@pytest.fixture(autouse=True)
+def mock_zmap() -> Generator[MagicMock, Any, None]:
+    """Automatically mock the ZMap instance before each test."""
+    with patch("zmapsdk.api.ZMap") as mock_zmap:
+        mock_instance = MagicMock()
+        mock_zmap.return_value = mock_instance
+        app.state.zmap = mock_instance
+        yield mock_instance
+        if hasattr(app.state, "zmap"):
+            del app.state.zmap
+
+
+@pytest.fixture(autouse=True)
+def setup_teardown() -> Generator[None, Any, None]:
+    """Fixture to handle setup and teardown for each test."""
+    mock_zmap = MagicMock()
+    app.state.zmap = mock_zmap
+    yield
+    # Teardown - clean up
+    if hasattr(app.state, "zmap"):
+        del app.state.zmap
+
+
+def test_root_endpoint_success(
+    client: TestClient, mock_zmap: Generator[MagicMock, Any, None]
+) -> None:
+    """Test the root endpoint with successful ZMap version check."""
+    mock_zmap.get_version.return_value = "zmap 3.0.0"
+
+    response = client.get("/")
+
+    assert response.status_code == 200
+    assert response.json() == {
+        "name": "ZMap SDK API",
+        "version": "zmap 3.0.0",
+        "description": "REST API for ZMap network scanner",
+    }
+    mock_zmap.get_version.assert_called_once()
+
+
+def test_api_docs_endpoints(client: TestClient) -> None:
+    """Test that api documentation endpoints exist."""
+    response = client.get("/openapi.json")
+    assert response.status_code == 200
+    assert "openapi" in response.json()
+
+    response = client.get("/docs")
+    assert response.status_code == 200
+    assert "swagger-ui" in response.text
+
+    response = client.get("/redoc")
+    assert response.status_code == 200
+    assert "redoc" in response.text

tests/test_config.py

@@ -0,0 +1,24 @@
+import pytest
+
+from zmapsdk import ZMapConfigError, ZMapScanConfig
+
+
+def test_config() -> None:
+    """Test config."""
+    conf = ZMapScanConfig(
+        target_port=80,
+        bandwidth=None,
+        rate=50,
+        max_targets=10,
+        min_hitrate=10,
+        vpn=False,
+    )
+    assert conf.target_port == 80
+    assert conf.to_dict().get("rate") == 50
+
+    with pytest.raises(ZMapConfigError) as exc_info:
+        conf = ZMapScanConfig(source_port=-9999)
+        assert (
+            f"Invalid source port range: {conf.source_port}. Must be between 0 and 65535."
+            in str(exc_info.value)
+        )

tests/test_schemas.py

@@ -0,0 +1,28 @@
+from zmapsdk.schemas import ScanRequest
+
+
+def test_minimal_valid_request():
+    """Test that minimal valid request works"""
+    request = ScanRequest()
+    assert request.target_port is None
+    assert request.return_results is False
+
+
+def test_full_valid_request():
+    """Test all fields with valid values"""
+    request = ScanRequest(
+        target_port=80,
+        subnets=["192.168.1.0/24", "10.0.0.0/8"],
+        output_file="/path/to/output.json",
+        blocklist_file="/path/to/blocklist.txt",
+        allowlist_file="/path/to/allowlist.txt",
+        bandwidth="10Mbps",
+        probe_module="tcp_syn",
+        rate=1000,
+        seed=42,
+        verbosity=2,
+        return_results=True,
+    )
+    assert request.target_port == 80
+    assert request.subnets == ["192.168.1.0/24", "10.0.0.0/8"]
+    assert request.return_results is True

zmapsdk/api.py

@@ -5,8 +5,8 @@
 import psutil
 import uvicorn
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse as StarletteFileResponse
 
+from zmapsdk.core import ZMap
 from zmapsdk.schemas import (
     BlocklistRequest,
     FileResponse,
@@ -15,8 +15,6 @@
     StandardBlocklistRequest,
 )
 
-from .core import ZMap
-
 # Scan tracking dictionary
 active_scans = {}