Writing tests and securing code with AI

AI-Assisted Coding for Developers

Francesca Donadoni

AI Curriculum Manager, DataCamp

Test-driven prompting

☺ Normal Prompting

Write a Python function to parse email addresses.

🤓 Test-Driven Prompting

Write a Python function to parse email addresses. It should pass the following tests:

Accepts a valid address like "[email protected]".

Rejects invalid input like "user@@domain".

Raises a ValueError for empty input.

Test-driven prompting

User prompt: Generate unit tests in Python for a function that processes user-submitted form data and inserts it into a SQL database.
Model response:

Here's a ready-to-run pytest suite that covers validation, normalization, SQL-injection safety, and transaction behavior for a function that processes form data and inserts a record into a SQL database.

[...]
# ---------- Test DB fixtures ----------

@pytest.fixture()
def conn():
    """In-memory SQLite DB with a users table."""
    c = sqlite3.connect(":memory:") [...]

Test-driven prompting

User prompt: Generate Python unit tests using pytest for the following function. Include cases for empty input, SQL keywords in the input, and special characters. Assume a mock database connection. <INSERT FUNCTION>
Model response:

Here's a self-contained pytest test module that uses mocks (no real DB) and covers empty input, SQL keywords, and special characters. It also verifies commit/close behavior and that exceptions are swallowed.

[...]
@pytest.fixture
def mock_sqlite_connect(monkeypatch):
    """
    Patch sqlite3.connect to return a MagicMock connection + cursor.
[...]

Prompting for security

☺ Normal Prompting

Scan the following Python function for potential vulnerabilities and suggest safer alternatives if any are found.

🤓 Test-Driven Prompting

Review this code for common web security issues, including SQL injection, XSS, and input validation flaws.

OWASP Top 10

User prompt: Audit the following function for the OWASP Top 10 vulnerabilities.

def login(username, password):
   query = f"SELECT * FROM users WHERE username = '{username}' \
           AND password = '{password}'"
   result = db.execute(query)
   return result.fetchone()

Model response:

Here's a focused OWASP Top 10 audit of your login function, plus a safer rewrite.

A03: Injection. Direct string interpolation into SQL makes this trivially SQL-injectable. Example payloads:
- username = "admin' --" → bypasses password check. [...]

Let's practice!

AI-Assisted Coding for Developers