Unit Test First Prompting: Generate Tests Before Implementation

Posted 12 Sep by JAMIUL ISLAM 0 Comments

Unit Test First Prompting: Generate Tests Before Implementation

You’ve probably been there. You ask an AI to write a function. It spits out code that looks perfect. You run it. It breaks. Or worse, it works for the happy path but crashes when someone enters a weird character in a username field. This is the trap of AI-assisted development: models are great at guessing plausible code, but terrible at knowing exactly what you meant until you show them.

Unit Test First Prompting flips this script. Instead of asking the AI to build the feature, you ask it to define the rules of the game first. You generate the tests before the implementation exists. This isn’t just a clever trick; it’s a shift in how we handle security and correctness in the age of large language models. By forcing the AI to create failing tests first, you turn vague requirements into executable specifications.

Why Writing Code First Fails With AI

When you prompt an LLM like GitHub Copilot or ChatGPT to "write a user validator," it makes assumptions. Does it allow special characters? What about spaces? Is it case-sensitive? The model fills these gaps with its training data, which might not match your specific business logic or security needs.

This leads to two major problems:

  • The Illusion of Correctness: The code compiles and runs, so you assume it’s right. But without explicit constraints, subtle bugs hide in edge cases.
  • Security Blind Spots: AI often ignores Common Weakness Enumeration (CWE) mitigations unless explicitly told. A generated input handler might look clean but fail to prevent SQL injection or buffer overflows because the prompt didn't demand it.

By starting with tests, you remove the guesswork. The tests become the contract. If the implementation doesn’t pass the tests, it’s wrong. Period. No arguing with the AI about whether a space should be allowed in a username. The test says no, so it’s no.

The Red-Green-Refactor Cycle for AI

If you’re familiar with Test-Driven Development (TDD), you know the drill: Red, Green, Refactor. Applying this to AI prompting requires adapting each stage to interact with a language model rather than writing code manually.

Stage 1: The Red Phase (Generate Tests)

In traditional TDD, you write a test that fails because the function doesn’t exist yet. In Unit Test First Prompting, you prompt the AI to generate only the test suite. Crucially, you instruct it to ignore implementation details.

A weak prompt looks like this: "Write tests for a username validator."

A strong prompt looks like this:

"Act as a senior QA engineer. I need unit tests for a new function called `validateUsername`. Requirements: Usernames must be 3-16 characters long, start with a letter, and contain only alphanumeric characters. Specifically include tests for CWE-20 (Improper Input Validation). Do not generate the implementation code yet. Just output the test functions."

This forces the AI to think about edge cases-empty strings, maximum length, invalid starting characters-before it worries about syntax. These tests will initially fail (or not compile) because the function is missing. That’s the "Red" state.

Stage 2: The Green Phase (Generate Implementation)

Now you feed those tests back to the AI along with the original requirement. You ask it to write the code that makes the tests pass. Because the tests are already defined, the AI has a clear target. It can’t hallucinate a different behavior if the test explicitly checks for one.

For example, if the test expects an exception for a 17-character username, the implementation must throw that exception. If the AI tries to truncate the string instead, the test fails, and you know immediately.

Stage 3: The Refactor Phase (Optimize)

Once the tests pass, you ask the AI to refactor the code for readability or performance, ensuring the tests still pass. This step is safer now because you have a safety net. If the refactored code breaks a test, you revert. Without tests, refactoring AI-generated code is risky-you might break functionality you didn’t realize existed.

Advanced Prompting Techniques for Better Tests

Not all prompts are created equal. Research shows that the structure of your request significantly impacts the quality of generated tests. Here are three techniques that consistently yield better results.

Role Priming

Start by assigning a persona. "Act as an enterprise security auditor" yields different tests than "Act as a junior developer." The former will aggressively probe for vulnerabilities; the latter might stick to basic functionality. For critical systems, always prime for rigor.

Few-Shot Prompting

Provide examples. Show the AI one or two existing test cases from your codebase. This helps the model understand your preferred testing framework (e.g., Jest vs. Pytest), naming conventions, and assertion styles. It reduces the need for manual cleanup later.

Scenario Enumeration

Don’t just ask for tests; list the scenarios. Explicitly mention normal operations, boundary conditions (min/max values), and error states. AI models struggle with implicit context. If you don’t say "test for null inputs," it might skip them.

Comparison of Prompting Strategies for Unit Test Generation
Strategy Description Best Use Case Risk Level
Simple Prompting Single instruction, e.g., "Write tests for X." Quick prototypes, simple utilities High (misses edge cases)
Chain-of-Thought Step-by-step reasoning: Analyze -> List Scenarios -> Generate Complex logic, algorithmic functions Low (high coverage)
Constraint-Based Explicit inclusion of CWEs and business rules Security-critical modules, public APIs Very Low (secure by design)
Robot shielded by green energy against attacking drones, symbolizing passing tests.

Integrating Security Into the Workflow

One of the biggest wins of Unit Test First Prompting is security integration. When you include CWE references in your initial test generation prompt, you bake security into the specification.

Consider a file upload feature. A standard prompt might generate code that accepts any file type. A security-focused test-first prompt would require tests that reject `.exe` files, check for path traversal attacks (`../`), and verify MIME types. If the implementation passes these tests, it likely handles these threats correctly. If it fails, you catch the vulnerability before deployment.

This approach shifts security from a post-development audit to a pre-development requirement. You’re not checking if the code is secure; you’re building it to satisfy security tests from day one.

Tools and Frameworks to Scale the Practice

Doing this manually for every function gets tedious. Fortunately, several tools support this workflow natively or through configuration.

GitHub Copilot offers direct integration. You can highlight code, right-click, and select "Generate Tests." However, for true test-first prompting, you often need to use Copilot Chat to iteratively refine the test suite before generating the code. Use slash commands like `/tests` to trigger generation, but follow up with specific prompts to add security cases.

For larger teams, consider using rule-based frameworks like .cursorrules or Markdown Configuration (.mdc) files. These act as "always-on" guardrails. You can define global instructions such as "Always generate tests before implementation" or "Include CWE-89 mitigation in all database query tests." This ensures consistency across developers without relying on individual discipline.

Polished giant robot standing in sunlight, representing refactored, stable code.

Common Pitfalls and How to Avoid Them

Even with a good strategy, things go wrong. Here’s what to watch out for.

  • Hallucinated Mocks: AI often invents mock objects that don’t exist in your library. Always review generated mocks against your actual dependencies.
  • Overloading Prompts: Don’t paste 500 lines of context into one prompt. Break it down. Generate tests for one function at a time. Too much noise degrades output quality.
  • Ignoring Compilation Errors: Sometimes the generated tests won’t even compile due to syntax errors or missing imports. Treat these as immediate feedback. Fix the prompt or the test code before moving to implementation.
  • False Positives: Ensure your tests actually fail when they should. Run them against a dummy implementation that returns hardcoded values. If they pass, your tests are too loose.

Practical Example: Building a Payment Processor

Let’s walk through a real-world scenario. You need a function `processPayment(amount, currency)`.

Step 1: Define Criteria. Amount must be positive. Currency must be USD, EUR, or GBP. Reject amounts over $10,000.

Step 2: Prompt for Tests.

"Generate pytest unit tests for `processPayment`. Test valid USD, EUR, GBP. Test negative amounts, zero amount, and unsupported currencies. Test the $10,000 limit boundary. Include assertions for expected return values."

Step 3: Review Tests. Did it test the boundary exactly at 10,000? Did it test floating-point precision issues? Add missing cases manually if needed.

Step 4: Prompt for Implementation.

"Write the Python implementation for `processPayment` that passes the following tests: [paste tests]. Use Decimal for monetary calculations to avoid float errors."

Step 5: Verify. Run the tests. If they pass, you’re done. If not, iterate.

Is It Worth the Effort?

Yes, especially if you’re working on production-grade software. While writing tests first feels slower upfront, it saves massive amounts of time debugging vague AI outputs later. It also creates documentation that stays current. If the tests change, the behavior changes. Your codebase becomes self-verifying.

For beginners, start small. Pick one utility function. Write the tests via AI. Then write the code. Notice how the clarity improves. Once you see the difference, you’ll never go back to "prompt and pray."

Do I need to know TDD to use Unit Test First Prompting?

No, but understanding the basics helps. You don’t need to write complex test frameworks manually; the AI does that. However, you do need to clearly articulate what the code should do. If you can describe requirements, you can drive this process.

Which AI models work best for this method?

Models with strong reasoning capabilities perform best. GPT-4o, Claude 3.5 Sonnet, and GitHub Copilot (powered by OpenAI models) are currently top performers for generating structured, logical test suites. Smaller models may miss edge cases or security nuances.

How do I handle legacy code with no tests?

You can reverse-engineer tests. Ask the AI to analyze the existing function and generate tests that reflect its current behavior. Run these tests to establish a baseline. Then, when modifying the code, ensure the new implementation passes the old tests plus any new ones you add for changed behavior.

Does this slow down development?

Initially, yes. But once you get used to the prompting patterns, it speeds up. You spend less time debugging unexpected behaviors and more time shipping features. The iterative nature prevents rework loops that plague traditional AI coding.

Can I automate this entire workflow?

Partially. Tools like Cursor or specialized IDE plugins can enforce test-first rules. However, human oversight is still critical for defining business logic and reviewing security implications. Full automation risks encoding incorrect assumptions into the test suite itself.

Write a comment