Testing your automation with pytest

Confidence to change scripts without breaking prod.

You change one line in a script that pulls findings from your vulnerability scanner and files them as tickets. On your laptop it looks right, so you push it. An hour later, every finding your scanner marked CRITICAL is landing in the queue as 'low', and nobody notices until the audit. People get nervous touching working automation because they have no fast, honest way to check they did not break something. A test is that check.

A test is a small saved experiment. You feed your code a known input, write down the answer you expect, and let a tool confirm that reality matches. It runs in a fraction of a second, as often as you like. pytest (a Python testing tool, said 'pie-test') is the standard way to write and run those experiments. This lesson builds a small security helper and wraps it in tests, including the part that trips people up in automation: tests that never touch a real scanner, a real server, or the real network.

A test is a file, a function, and one assert

pytest finds tests by convention. It looks for files named test_*.py, and inside them, functions whose names start with test_. Inside a function you state a claim with Python's built-in assert keyword. Writing assert count == 2 means 'I expect count to be 2; stop and complain loudly if it is not.' There are no special assertion methods to memorize. Install the tool with pip (Python's package installer): pip install pytest.

def parse_findings(text):
    """Turn scanner text into a list of findings.
 
    Each real line looks like: 'CVE-2024-3094 CRITICAL xz-utils'.
    Blank lines and comment lines (starting with #) are ignored.
    """
    findings = []
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split()
        if len(parts) < 3:
            continue
        cve, severity, package = parts[0], parts[1], parts[2]
        findings.append({"cve": cve, "severity": severity.upper(), "package": package})
    return findings
 
 
def high_severity_count(findings):
    return sum(1 for f in findings if f["severity"] in ("HIGH", "CRITICAL"))

parse_findings turns the scanner's text output into a tidy list of dictionaries, skipping blank lines and comment lines. high_severity_count answers one question: how many findings are HIGH or CRITICAL? Both are pure functions, meaning the same input always gives the same output with no side effects. That makes them the easiest kind of code to test.

import scanner
 
 
def test_parse_findings_basic():
    text = "CVE-2024-3094 CRITICAL xz-utils"
    result = scanner.parse_findings(text)
    assert result == [
        {"cve": "CVE-2024-3094", "severity": "CRITICAL", "package": "xz-utils"}
    ]
 
 
def test_parse_findings_skips_comments_and_blanks():
    text = "# scan results\n\nCVE-2021-44228 high log4j\n"
    result = scanner.parse_findings(text)
    assert len(result) == 1
    assert result[0]["severity"] == "HIGH"

Save both files in the same folder and run pytest there. It discovers the file, runs both functions, and reports back.

Each dot is a passing test. The [100%] tracks progress through the run. The green summary line is the whole point: two experiments ran, both matched. The real value shows up the moment your code is wrong. Suppose an edit drops the .upper() call, so a scanner line reading 'high' no longer becomes 'HIGH'. Rerun the same command.

pytest turned your plain assert into a bug report. It points at the exact line, shows what it got ('high') against what you expected ('HIGH'), and even draws the difference. That gap matters for security: a downstream filter checking severity == 'HIGH' would silently skip a real high-severity finding. The test caught it in under two seconds, before it ever shipped.

Fixtures: setup you write once

A fixture is like the tray of prepped ingredients a cook lays out before service: measured, ready, identical for every dish. In pytest, a fixture is a function that builds some setup or test data. Any test that lists the fixture's name as an argument gets a fresh copy handed to it. Fresh is the point, because one test can then never pollute another.

import pytest
 
 
@pytest.fixture
def sample_report():
    return (
        "# scan results\n"
        "CVE-2024-3094 CRITICAL xz-utils\n"
        "CVE-2021-44228 HIGH log4j\n"
        "CVE-2023-0001 LOW zlib\n"
    )
 
 
def test_high_severity_count(sample_report):
    findings = scanner.parse_findings(sample_report)
    count = scanner.high_severity_count(findings)
    assert count == 2

pytest sees the sample_report parameter, runs the fixture with the matching name, and passes its return value in. Write that report once and ten tests can share it, each with its own clean copy.

Many cases from one test, with parametrize

Instead of photocopying a test ten times and changing one box on each copy, you hand over a list of values and let the machine stamp out every case. The @pytest.mark.parametrize decorator (a label that means 'run this test once per row') does exactly that. You give it pairs of input and expected answer; pytest runs the body once for each pair.

@pytest.mark.parametrize(
    "line, expected",
    [
        ("CVE-1 critical pkg", "CRITICAL"),
        ("CVE-2 High pkg", "HIGH"),
        ("CVE-3 low pkg", "LOW"),
    ],
)
def test_severity_is_uppercased(line, expected):
    findings = scanner.parse_findings(line)
    assert findings[0]["severity"] == expected

The -v flag (verbose) prints one line per case, with the inputs baked into the test's name. Each row is now its own test. When one input breaks, you see precisely which one, not a vague 'something failed'.

The real skill: faking the outside world

A flight simulator lets a pilot rehearse an engine fire without setting a real plane alight. Mocking is the same move for code. You replace a real call to something slow, external, or dangerous with a stand-in you fully control. Your code runs unchanged; the thing it reaches for becomes a puppet.

Automation reaches outward constantly. run_scan launches a scanner through subprocess (Python's built-in way of starting another program and reading what it prints). fetch_cve_score looks up a CVE (Common Vulnerabilities and Exposures, the public catalog that gives every known security bug a unique ID like CVE-2024-3094) by calling an API (application programming interface, a web service you send requests to) with the requests library. You do not want a test launching a real scanner or phoning a real server: that is slow, flaky, needs credentials, and can change real state. So you replace those two calls with fakes. Two tools do it: monkeypatch, a fixture pytest hands you, and unittest.mock, part of Python's standard library.

import subprocess
 
import requests
 
 
def run_scan(target):
    """Run the scanner CLI and return parsed findings.
 
    Raises RuntimeError if the scanner exits non-zero.
    """
    result = subprocess.run(
        ["trivy", "image", "--quiet", target],
        capture_output=True,
        text=True,
        timeout=60,
    )
    if result.returncode != 0:
        raise RuntimeError(f"scan failed: {result.stderr.strip()}")
    return parse_findings(result.stdout)
 
 
def fetch_cve_score(cve):
    """Look up a CVE's CVSS base score from an API."""
    resp = requests.get(f"https://api.example.com/cve/{cve}", timeout=10)
    resp.raise_for_status()
    return resp.json()["score"]

run_scan builds a command, runs it with a 60-second timeout, and treats a non-zero exit code as failure. A non-zero exit code is how a Unix program signals 'something went wrong'. CompletedProcess is the object subprocess.run normally returns once the program finishes.

import subprocess
 
 
def test_run_scan_parses_output(monkeypatch):
    fake = subprocess.CompletedProcess(
        args=["trivy"],
        returncode=0,
        stdout="CVE-2024-3094 CRITICAL xz-utils\n",
        stderr="",
    )
    monkeypatch.setattr(subprocess, "run", lambda *a, **k: fake)
    findings = scanner.run_scan("nginx:latest")
    assert findings[0]["cve"] == "CVE-2024-3094"

monkeypatch.setattr(subprocess, 'run', ...) replaces the real function with your lambda (a one-line throwaway function) for the length of this one test, then pytest quietly restores the original afterward. Your fake returns a CompletedProcess you built by hand, so run_scan parses your canned output and launches nothing at all.

Test the unhappy paths on purpose

A fire drill exists so you find the locked exit on a calm Tuesday, not during a real fire. The failure paths are where automation hurts you in production (the live systems your team depends on): a scanner that dies mid-run, an API that hangs forever. Test them on purpose.

def test_run_scan_raises_on_nonzero_exit(monkeypatch):
    fake = subprocess.CompletedProcess(
        args=["trivy"],
        returncode=1,
        stdout="",
        stderr="db download failed",
    )
    monkeypatch.setattr(subprocess, "run", lambda *a, **k: fake)
    with pytest.raises(RuntimeError, match="scan failed"):
        scanner.run_scan("nginx:latest")
 
 
def test_run_scan_propagates_timeout(monkeypatch):
    def boom(*a, **k):
        raise subprocess.TimeoutExpired(cmd="trivy", timeout=60)
 
    monkeypatch.setattr(subprocess, "run", boom)
    with pytest.raises(subprocess.TimeoutExpired):
        scanner.run_scan("nginx:latest")

pytest.raises(RuntimeError) is a context manager meaning 'the code in this block had better raise that exception; fail if it does not'. For the non-zero exit, you hand run_scan a CompletedProcess with returncode set to 1 and confirm it raises RuntimeError. For the timeout, your fake raises subprocess.TimeoutExpired, the very exception the real call throws after 60 seconds, and you check that run_scan lets it travel up. You proved the error handling works in milliseconds, without waiting a minute or breaking a live scan.

from unittest.mock import patch, MagicMock
 
 
def test_fetch_cve_score():
    fake_resp = MagicMock()
    fake_resp.json.return_value = {"score": 10.0}
    fake_resp.raise_for_status.return_value = None
    with patch("scanner.requests.get", return_value=fake_resp) as mock_get:
        score = scanner.fetch_cve_score("CVE-2024-3094")
    assert score == 10.0
    mock_get.assert_called_once()

For the API, unittest.mock.patch swaps scanner.requests.get for a MagicMock, a fake object that conjures whatever attributes and methods you ask of it. You tell its .json() what to return, then assert your code pulled the right CVSS (Common Vulnerability Scoring System) score, a 0-to-10 severity number, out of the response. mock_get.assert_called_once() then confirms you actually called the API exactly once, so a typo that skips the call still fails the test.

The -k flag (keyword) runs only tests whose names match your expression. Four green: the happy scanner path, both failure paths, and the API call, every one verified without touching a real system.

Running it in CI

A turnstile lets you through only after you tap a valid card. CI (continuous integration, a service that runs your checks automatically on every push) is that turnstile for code. You wire pytest into it so no change reaches the main branch until every test passes. Here is a small GitHub Actions workflow that does it.

name: tests
on: [push, pull_request]
 
jobs:
  pytest:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
      - run: pip install pytest requests
      - run: pytest -q

On every push and pull request, it checks out your code, sets up Python 3.11, installs your dependencies, and runs pytest -q (quiet mode, one line of summary). If any test fails, that step exits non-zero and the merge is blocked. The gate is automatic, so nobody has to remember to run the tests.

For a security engineer, this is the gap between shipping a scanner change on a Friday afternoon and lying awake about it. Remember the edit that quietly turned every CRITICAL finding into 'low'? A three-line test would have caught it at your desk, in half a second, before it reached the ticket queue. Write the test once, and every change you make from then on has to earn its way past it.

Related