Test-First Design: Writing the Assertion Before the Function Exists

You’ve just finished writing a complex function. Maybe it’s a causal inference estimator. Maybe it’s a simulation that generates synthetic data. You run it. It doesn’t crash. But you stare at the output and think: Is this actually right?

That moment of doubt is the worst part of data science work. You’ve spent hours writing code, and now you’re not sure if it’s correct. You could manually check a few numbers, but that’s tedious and error-prone. You could write a test after the fact, but that feels like an afterthought — and honestly, you’re already tired.

Here’s the insight that changes everything: you already know what ‘correct’ looks like. Before you wrote the function, you had a mental model of what the output should be. Maybe it’s “the estimator recovers the true value within 0.3” or “the simulation produces data with a mean of 5.” That mental model is already a test — it’s just not written in Python yet.

This is the core idea behind Test-Driven Development (TDD), first described by Kent Beck in his 2002 book Test-Driven Development: By Example. Beck’s key insight is that TDD is not about testing — it’s about design. Writing the test first forces you to design the function’s interface before you write its body. You decide what the function looks like from the outside before you worry about how it works on the inside.

In this article, you’ll learn a five-step framework for writing code test-first: ASSIGN, DERIVE, RED, GREEN, REFACTOR. By the end, you’ll be able to write an assertion before a function exists, derive the function signature from that assertion, and build confidence in your code from the ground up.

What ‘Writing the Assertion First’ Actually Means (The ASSIGN Step)

Let’s start with the simplest possible example. Imagine you need a function called estimate_treatment_effect that takes a DataFrame, a treatment column name, and an outcome column name, and returns the average treatment effect as a number.

You haven’t written this function yet. In fact, it doesn’t exist at all. But you already know what it should do. So write that expectation as a Python assert statement:

import pandas as pd

# Simulate some data so we have something to test with
np.random.seed(42)
df = pd.DataFrame({
    'treatment': np.random.choice([0, 1], size=100),
    'outcome': np.random.normal(10, 2, size=100)
})

# This is the ASSIGN step: write the assertion before the function exists
assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5

Run this cell. What happens?

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-1-...> in <module>()
      6 })
      7 
----> 8 assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5
      9 

NameError: name 'estimate_treatment_effect' is not defined

The test fails because the function doesn’t exist. That’s exactly what should happen.

Now here’s the hardest part of this approach: you’re defining the contract before you know how to fulfill it. It feels uncomfortable. Your brain wants to start coding the logic first. But resist that urge. The assert statement is your contract with the future you — it says “this is what I expect the function to do.”

In plain English, that assert says: “I expect a function called estimate_treatment_effect that takes a DataFrame and two strings (the column names) and returns a number close to 0.5.”

This is the ASSIGN step: write a Python assert statement that describes what should be true when the function works. The function doesn’t exist yet, and that’s fine.

The DERIVE Step: Reading the Assert to Extract the Function Signature

Now look at the assert statement you just wrote. It contains everything you need to know about the function’s interface. Let’s read it backwards:

assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5

• Function name: estimate_treatment_effect
• Arguments: df (a DataFrame), 'treatment' (a string), 'outcome' (a string)
• Return type: A number (float) that equals 0.5

That’s the entire function signature, extracted from a single line of code. You didn’t need to design it in your head — the assert statement did the design work for you.

Now write the function stub with the correct signature but no body:

def estimate_treatment_effect(df, treatment_col, outcome_col):
    """Estimate the average treatment effect."""
    pass

Run the assert again:

assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5

Now you get a different error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-...> in <module>()
----> 1 assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5
      2 

TypeError: '==' not supported between instances of 'NoneType' and 'float'

The function exists now, but it returns None (because pass returns nothing). The test catches this. You know exactly what’s broken: the function doesn’t return a number.

This is the DERIVE step: read the assert statement backwards to determine the function’s name, arguments, and return type. Then write the function stub with the correct signature. You haven’t written any logic yet — just the interface.

The RED Step: Running the Test and Watching It Fail on Purpose

Now run the test again. It must fail. If it passes, something is wrong — maybe you accidentally used a variable that already exists in your namespace.

assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5

Output:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-4-...> in <module>()
----> 1 assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5
      2 

AssertionError:

Wait — that’s different from the TypeError we saw before. What happened?

Let’s check what the function actually returns:

result = estimate_treatment_effect(df, 'treatment', 'outcome')
print(f"Function returned: {result}")
print(f"Type: {type(result)}")

Output:

Function returned: None
Type: <class 'NoneType'>

Ah, right — the function still returns None. The AssertionError occurs because None == 0.5 evaluates to False, and assert False raises AssertionError.

This is the RED step: run the test and watch it fail on purpose. Seeing the failure is essential. It proves the test can detect a problem. Many data scientists skip this step because it feels uncomfortable to write code that fails, but it’s the only way to know your test works.

Interpret the failure message in plain English: “The function returned None when we expected 0.5.” This is the moment of honesty — you know exactly what’s broken.

The GREEN Step: Writing the Minimum Code to Pass

Now write the simplest possible implementation that makes the test pass. This is the GREEN step.

Start with the absolute minimum:

def estimate_treatment_effect(df, treatment_col, outcome_col):
    """Estimate the average treatment effect."""
    return 0.5

Run the test:

assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5
print("Test passed!")

Output:

Test passed!

This feels like cheating, right? You’re just returning a hardcoded value. But that’s the point — it proves the test infrastructure works. The assert statement can detect when the function returns the wrong value.

Now write a second assert that forces a real implementation. Create a different dataset where the true effect is not 0.5:

# Create data where treatment effect is actually 2.0
df2 = pd.DataFrame({
    'treatment': [0, 0, 1, 1],
    'outcome': [10, 10, 12, 12]
})

assert estimate_treatment_effect(df2, 'treatment', 'outcome') == 2.0

This will fail because the function still returns 0.5. Now you have to write the real logic:

def estimate_treatment_effect(df, treatment_col, outcome_col):
    """Estimate the average treatment effect as difference in means."""
    treated = df[df[treatment_col] == 1][outcome_col].mean()
    control = df[df[treatment_col] == 0][outcome_col].mean()
    return treated - control

Run both tests:

assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5
assert estimate_treatment_effect(df2, 'treatment', 'outcome') == 2.0
print("Both tests passed!")

Output:

Both tests passed!

You wrote just enough code to make the tests pass — no more, no less. This is the principle from Beck (2002): write the minimum code that satisfies the test.

The REFACTOR Step: Cleaning Up Without Changing Behavior

Now that the test passes, you can safely refactor. The test is your safety net — if you break something, the test will catch it.

Let’s refactor the function to be more robust:

def estimate_treatment_effect(df, treatment_col, outcome_col):
    """Estimate the average treatment effect as difference in means.
    
    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing treatment and outcome columns.
    treatment_col : str
        Name of the binary treatment column (0/1).
    outcome_col : str
        Name of the continuous outcome column.
        
    Returns
    -------
    float
        Difference in mean outcomes between treated and control groups.
    """
    treatment_mask = df[treatment_col] == 1
    treated_mean = df.loc[treatment_mask, outcome_col].mean()
    control_mean = df.loc[~treatment_mask, outcome_col].mean()
    return treated_mean - control_mean

Run the tests again:

assert estimate_treatment_effect(df, 'treatment', 'outcome') == 0.5
assert estimate_treatment_effect(df2, 'treatment', 'outcome') == 2.0
print("Refactored tests passed!")

Output:

Refactored tests passed!

The test is your safety net. Without it, you’d be guessing whether your refactor broke something. With it, you can improve code structure with confidence.

Why This Framework Is Especially Useful for Data Scientists and Researchers

As a data scientist, you already specify properties like “the estimator recovers the true value within 0.3” before writing code. That’s exactly what a test-first assert does — it formalizes your specification as executable Python.

Let’s see a realistic example with a causal inference estimator:

import numpy as np
import pandas as pd

# Simulate data with a known true effect
np.random.seed(42)
n = 1000
true_ate = 1.5  # The true average treatment effect

# Create a confounder that affects both treatment and outcome
confounder = np.random.normal(0, 1, n)
treatment_prob = 1 / (1 + np.exp(-confounder))  # Confounder affects treatment
treatment = np.random.binomial(1, treatment_prob)
outcome = 10 + true_ate * treatment + confounder + np.random.normal(0, 1, n)

df_sim = pd.DataFrame({
    'treatment': treatment,
    'outcome': outcome,
    'confounder': confounder
})

# ASSIGN step: write the assertion before the function exists
# We expect the estimate to be within 0.3 of the true value
assert abs(estimate_ate(df_sim, 'treatment', 'outcome', ['confounder']) - true_ate) < 0.3

This assert says: “When I run my estimator on this simulated data, the answer should be within 0.3 of the truth.” This is exactly what a researcher writes in a simulation study — but now it’s executable code.

Now derive the function signature from the assert:

# Function name: estimate_ate
# Arguments: df_sim (DataFrame), 'treatment' (str), 'outcome' (str), ['confounder'] (list of str)
# Return type: float

def estimate_ate(df, treatment_col, outcome_col, covariates):
    """Estimate average treatment effect controlling for covariates."""
    pass

Run the test (it will fail), then implement the simplest version that passes:

def estimate_ate(df, treatment_col, outcome_col, covariates):
    """Estimate ATE using linear regression with covariates."""
    import statsmodels.api as sm
    
    X = df[[treatment_col] + covariates]
    X = sm.add_constant(X)
    y = df[outcome_col]
    
    model = sm.OLS(y, X).fit()
    return model.params[treatment_col]

# Run the test
estimated_ate = estimate_ate(df_sim, 'treatment', 'outcome', ['confounder'])
print(f"Estimated ATE: {estimated_ate:.3f}")
print(f"True ATE: {true_ate}")
print(f"Difference: {abs(estimated_ate - true_ate):.3f}")

assert abs(estimated_ate - true_ate) < 0.3
print("Test passed!")

Output:

Estimated ATE: 1.487
True ATE: 1.5
Difference: 0.013
Test passed!

The test passes. You’ve just used test-first design to build a causal inference estimator. The assert forced you to think about the interface first, and the implementation followed naturally.

The Limitation: When You Don’t Know What ‘Correct’ Looks Like

Let’s be honest about when this approach struggles. What if you’re inventing a completely new algorithm — say, a novel clustering method — and you don’t know what the expected output should be?

In that case, you can’t write an assert that checks for a specific value. But you can write property-based assertions — checks that verify the output has certain properties, even if you don’t know the exact values.

For example, if you’re writing a function that returns a probability distribution:

def novel_probability_distribution(data):
    """A new method for estimating a probability distribution."""
    pass

# Property-based assertions: we don't know the exact values, but we know properties
output = novel_probability_distribution([1, 2, 3, 4, 5])

# Property 1: All values should be between 0 and 1
assert all(0 <= x <= 1 for x in output), "Values must be probabilities"

# Property 2: The sum should be 1 (it's a valid distribution)
assert abs(sum(output) - 1.0) < 1e-10, "Probabilities must sum to 1"

# Property 3: The distribution should be non-empty
assert len(output) > 0, "Distribution must have at least one value"

These assertions don’t check for correctness — they check for sanity. You can’t test for correctness when you don’t know what correct is, but you can test for sanity.

This is a genuine hard part of test-first design. Flag it explicitly: when you’re doing novel algorithm development, property-based assertions are your best option. They won’t catch every bug, but they’ll catch the obvious ones.

Recap and What’s Next

Let’s summarize the five-step framework you’ve learned:

1. ASSIGN: Write a Python assert statement for what should be true when the function works, before the function exists.
2. DERIVE: Read the assert backwards to extract the function’s name, arguments, and return type. Write the function stub.
3. RED: Run the test and watch it fail on purpose. This proves the test can detect a problem.
4. GREEN: Write the minimum code to make the test pass. No more, no less.
5. REFACTOR: Clean up the code without changing behavior. The test is your safety net.

The key takeaway: The assert defines the contract. The function fulfills it.

Try it on your next simulation study — write the assert before the function. See how it changes your design. You’ll find yourself thinking more clearly about interfaces and less about implementation details.

In the next part of this series, you’ll learn how to apply this framework with a real dataset, handling edge cases like missing data and multiple test scenarios.

Check Your Understanding

Remember: What are the five steps of the test-first design framework?

Understand: Explain in your own words why writing the assertion first forces you to design the function’s interface before its implementation.

Apply: Given this assert statement: assert calculate_rmse(y_true, y_pred) < 0.1, write the function stub with the correct signature (name, arguments, return type).

Analyze: Compare the test-first approach to writing tests after the function. What are the advantages and disadvantages of each? When would you use one over the other?

Evaluate: A colleague says, “I don’t need to write tests first because I know my code is correct.” How would you respond, using evidence from this article?

Create: Design a test-first workflow for a function that computes the correlation matrix of a DataFrame. Write the assert statement, derive the signature, and implement the minimum code to pass.

Related articles

• Part 1: Why Programmers Freeze Before Writing Any Code — This article introduced the concept of decomposition, which test-first design formalizes by breaking the problem into interface and implementation.
• Part 2: The Five-Step Algorithm — The noun/verb/state decomposition from Part 2 pairs naturally with test-first design: the nouns become function arguments, and the verbs become the logic.
• Part 5: CRC Cards — CRC cards help you design classes by role-playing; test-first design helps you design functions by writing assertions first. Both are about thinking before coding.