AI-Assisted Code Maintenance and Improvement: Leveraging Test-Driven Development for Nonstandard Frameworks

Introduction

In modern software development, AI-powered tools have become indispensable for code maintenance and improvement. However, these tools are primarily trained on publicly available code from repositories like GitHub and discussions from platforms like Stack Overflow. This creates a significant challenge for software engineers working in companies with private codebases built on nonstandard frameworks or custom libraries.

In such environments, AI tools often fail to fully understand the application logic, leading to inaccurate suggestions and incorrect modifications. A practical approach to overcoming this limitation is to integrate Test-Driven Development (TDD) into AI-assisted workflows. By generating unit tests before modifying code, AI can iteratively refine its changes until they pass all tests, ensuring accuracy and alignment with the intended behavior.

The Challenge of AI in Nonstandard Codebases

Why AI Struggles with Custom Frameworks

1. Lack of Training Data: AI models are trained on publicly available data, meaning private or proprietary frameworks remain unknown to them.
2. Hidden Dependencies: AI lacks visibility into internal configurations, middleware, and hidden dependencies within custom frameworks.
3. Context Gaps: Without clear documentation, AI tools struggle to infer relationships between modules, leading to incorrect assumptions about function behaviors.

For companies using proprietary technology, this makes AI-driven maintenance unreliable—unless an effective strategy is applied.

Test-Driven Development (TDD) as a Solution

TDD is a software development methodology where tests are written before the actual implementation. Applying TDD in AI-assisted workflows ensures that the AI:

• Gains a structured understanding of the function.
• Validates its modifications through test execution.
• Iterates until it achieves the correct logic.

Steps for AI-Assisted Code Modification with TDD

1. Identify the Function to Modify

Before making changes, pinpoint the specific function and gather contextual information about its purpose and dependencies.

2. Generate a Unit Test First

Instead of modifying the function immediately, instruct the AI to create unit tests that capture the expected behavior. This forces the AI to understand the function’s input-output relationships.

3. Run the Test Suite

Execute the generated test suite to verify existing behavior. If tests fail, AI can iteratively refine them until they accurately capture expected functionality.

4. Modify the Function Based on Failing Tests

With a well-defined test suite, AI can now attempt modifications. If changes cause test failures, AI iterates through revisions until all tests pass.

5. Validate and Deploy

Once all tests pass, validate the changes manually to ensure correctness before deploying.

Example Workflow in JavaScript

Original Function with Unclear Logic

// ./discount.js
function calculateDiscount(price, userType) {
    // Validate price input
    if (typeof price !== 'number') {
        throw new Error('Price must be a number');
    }

    // Validate user type
    if (!userType || !['premium', 'gold', 'regular'].includes(userType)) {
        throw new Error('Invalid user type');
    }

    if (userType === "premium") {
        return price * 0.8;  // 20% discount
    } else if (userType === "gold") {
        return price * 0.85; // 15% discount
    } else {
        return price * 0.9;  // 10% discount
    }
}

module.exports = { calculateDiscount };

Step 1: Generate Unit Tests

To guide AI, we provide a clear prompt:

Prompt:

Generate Jest unit tests for the function calculateDiscount. Cover standard cases for premium, gold, and regular users, as well as edge cases including price = 0, null or undefined user type, unrecognized user type, negative price values, and non-numeric price inputs. Use Jest for testing with strict assertions.

Expected AI Output:

// Jest unit tests for calculateDiscount function
const { calculateDiscount } = require('./discount');

describe('calculateDiscount', () => {
    // Standard cases
    test('should apply 20% discount for premium users', () => {
        expect(calculateDiscount(100, 'premium')).toBe(80);
    });

    test('should apply 15% discount for gold users', () => {
        expect(calculateDiscount(100, 'gold')).toBe(85);
    });

    test('should apply 10% discount for regular users', () => {
        expect(calculateDiscount(100, 'regular')).toBe(90);
    });

    // Edge cases - Price variations
    test('should handle zero price correctly', () => {
        expect(calculateDiscount(0, 'premium')).toBe(0);
        expect(calculateDiscount(0, 'gold')).toBe(0);
        expect(calculateDiscount(0, 'regular')).toBe(0);
    });

    test('should handle negative price values', () => {
        expect(calculateDiscount(-100, 'premium')).toBe(-80);
        expect(calculateDiscount(-100, 'gold')).toBe(-85);
        expect(calculateDiscount(-100, 'regular')).toBe(-90);
    });

    test('should handle non-numeric price inputs', () => {
        expect(() => calculateDiscount('100', 'premium')).toThrow('Price must be a number');
        expect(() => calculateDiscount(null, 'premium')).toThrow('Price must be a number');
        expect(() => calculateDiscount(undefined, 'premium')).toThrow('Price must be a number');
    });

    // Edge cases - User type variations
    test('should handle null or undefined user type', () => {
        expect(() => calculateDiscount(100, null)).toThrow('Invalid user type');
        expect(() => calculateDiscount(100, undefined)).toThrow('Invalid user type');
    });

    test('should handle unrecognized user type', () => {
        expect(() => calculateDiscount(100, 'platinum')).toThrow('Invalid user type');
        expect(() => calculateDiscount(100, '')).toThrow('Invalid user type');
    });

    // Decimal precision tests
    test('should handle decimal prices correctly', () => {
        expect(calculateDiscount(99.99, 'premium')).toBeCloseTo(79.992, 3);
        expect(calculateDiscount(99.99, 'gold')).toBeCloseTo(84.9915, 3);
        expect(calculateDiscount(99.99, 'regular')).toBeCloseTo(89.991, 3);
    });
}); 

Benefits of AI-Assisted TDD in Nonstandard Codebases

1. Enhances AI’s Contextual Understanding: Writing tests first ensures AI comprehends function behavior before modifying it.
2. Prevents Regression Bugs: By continuously running tests, AI avoids introducing breaking changes.
3. Increases AI Accuracy: AI iterates until tests succeed, leading to correct modifications.
4. Creates Implicit Documentation: Unit tests serve as documentation for future engineers and AI models.

References and Further Reading

• Test-Driven Development: By Example – Kent Beck
• Refactoring: Improving the Design of Existing Code – Martin Fowler
• How AI Can Improve the Software Testing Process
• The Impact of AI on Software Testing

Conclusion

Maintaining and improving a nonstandard codebase with AI requires a structured approach. By leveraging Test-Driven Development, engineers can provide AI with the necessary context to make accurate modifications. This method not only enhances AI reliability but also ensures maintainability and robustness in private, proprietary applications.

By integrating AI with TDD, companies using custom frameworks can overcome the limitations of AI models trained on public code, enabling efficient and accurate software maintenance.