feat: AI-powered test case generation

.github/workflows/go.yml

@@ -54,3 +54,81 @@ jobs:
         files: ./coverage.out
         fail_ci_if_error: false
         verbose: true
+
+  e2e-ai-test:
+    name: E2E AI Tests with Real Ollama
+    runs-on: ubuntu-latest
+    needs: test  # Run after unit tests pass
+    steps:
+
+    - name: Set up Go
+      uses: actions/setup-go@v5
+      with:
+        go-version: '1.25.x'
+
+    - name: Check out code
+      uses: actions/checkout@v4
+
+    - name: Install Ollama
+      run: |
+        echo "Installing Ollama..."
+        curl -fsSL https://ollama.com/install.sh | sh
+        echo "Ollama installed successfully"
+
+    - name: Start Ollama service in background
+      run: |
+        echo "Starting Ollama service..."
+        ollama serve > /tmp/ollama.log 2>&1 &
+        OLLAMA_PID=$!
+        echo "Ollama PID: $OLLAMA_PID"
+        echo "Waiting for Ollama to start..."
+        sleep 5
+
+        # Verify Ollama is running
+        if curl -f http://localhost:11434/api/tags; then
+          echo "✓ Ollama is running"
+        else
+          echo "✗ Ollama failed to start"
+          cat /tmp/ollama.log
+          exit 1
+        fi
+
+    - name: Pull qwen2.5-coder:0.5b model
+      run: |
+        echo "Pulling qwen2.5-coder:0.5b model (400MB, ~2-3 minutes)..."
+        ollama pull qwen2.5-coder:0.5b
+        echo "✓ Model downloaded successfully"
+
+    - name: Verify model is available
+      run: |
+        echo "Verifying qwen2.5-coder:0.5b model..."
+        if ollama list | grep -q "qwen2.5-coder:0.5b"; then
+          echo "✓ Model is available"
+          ollama list
+        else
+          echo "✗ Model not found"
+          ollama list
+          exit 1
+        fi
+
+    - name: Get Go dependencies
+      run: go mod download
+
+    - name: Run E2E AI Tests
+      run: |
+        echo "Running E2E tests with real Ollama + qwen2.5-coder:0.5b..."
+        echo "These tests validate that AI generation matches golden files"
+        go test -v -tags=e2e -timeout=15m -coverprofile=e2e-coverage.out -covermode=count ./internal/ai
+      env:
+        CI: "true"
+        GOTESTS_E2E: "true"
+
+    - name: Upload E2E coverage to Codecov
+      uses: codecov/codecov-action@v5
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        files: ./e2e-coverage.out
+        flags: e2e-tests
+        name: e2e-coverage
+        fail_ci_if_error: false
+        verbose: true

.gitignore

@@ -1,2 +1,4 @@
 .DS_Store
-.claude/
+.claude/settings.local.json
+coverage*
+gotests_bin

CLAUDE.md

@@ -99,3 +99,4 @@ Existing test functions are automatically excluded to avoid duplication.
 - Tests are in `testdata/` directories with golden file comparisons in `testdata/goldens/`
 - The `templates/` directory contains built-in template sets
 - Bindata is used to embed templates in the binary (via `internal/render/bindata/`)
+- Always use scripts/regenerate-goldens.sh to generate the goldens for tests.

README.md

@@ -73,9 +73,153 @@ Available options:
 
   -use_go_cmp           use cmp.Equal (google/go-cmp) instead of reflect.DeepEqual
 
+  -ai                   generate test cases using AI (requires Ollama)
+
+  -ai-model             AI model to use (default "qwen2.5-coder:0.5b")
+
+  -ai-endpoint          Ollama API endpoint (default "http://localhost:11434")
+
+  -ai-min-cases         minimum number of test cases to generate with AI (default 3)
+
+  -ai-max-cases         maximum number of test cases to generate with AI (default 10)
+
   -version              print version information and exit
 ```
 
+## AI-Powered Test Generation
+
+**gotests** can generate intelligent test cases using local LLMs via [Ollama](https://ollama.ai). This feature analyzes your function implementations and generates realistic test values, edge cases, and error conditions.
+
+### Quick Start
+
+1. **Install Ollama** ([https://ollama.ai](https://ollama.ai))
+
+2. **Pull a model:**
+   ```sh
+   ollama pull qwen2.5-coder:0.5b  # Small, fast model (400MB)
+   # or
+   ollama pull llama3.2:latest     # Larger, more capable (2GB)
+   ```
+
+3. **Generate tests with AI:**
+   ```sh
+   gotests -all -ai -w yourfile.go
+   ```
+
+### Example
+
+Given this function:
+```go
+func CalculateDiscount(price float64, percentage int) (float64, error) {
+    if price < 0 {
+        return 0, errors.New("price cannot be negative")
+    }
+    if percentage < 0 || percentage > 100 {
+        return 0, errors.New("percentage must be between 0 and 100")
+    }
+    discount := price * float64(percentage) / 100.0
+    return price - discount, nil
+}
+```
+
+The AI generates (showing 3 cases; by default, the AI generates between 3-10 cases):
+```go
+func TestCalculateDiscount(t *testing.T) {
+    type args struct {
+        price      float64
+        percentage int
+    }
+    tests := []struct {
+        name    string
+        args    args
+        want    float64
+        wantErr bool
+    }{
+        {
+            name: "valid discount",
+            args: args{price: 100.0, percentage: 20},
+            want: 80.0,
+            wantErr: false,
+        },
+        {
+            name: "negative price",
+            args: args{price: -10.0, percentage: 20},
+            want: 0,
+            wantErr: true,
+        },
+        {
+            name: "invalid percentage",
+            args: args{price: 100.0, percentage: 150},
+            want: 0,
+            wantErr: true,
+        },
+    }
+    for _, tt := range tests {
+        t.Run(tt.name, func(t *testing.T) {
+            got, err := CalculateDiscount(tt.args.price, tt.args.percentage)
+            if (err != nil) != tt.wantErr {
+                t.Errorf("CalculateDiscount() error = %v, wantErr %v", err, tt.wantErr)
+                return
+            }
+            if got != tt.want {
+                t.Errorf("CalculateDiscount() = %v, want %v", got, tt.want)
+            }
+        })
+    }
+}
+```
+
+### AI Options
+
+```sh
+# Use a different model
+gotests -all -ai -ai-model llama3.2:latest -w yourfile.go
+
+# Generate a specific number of test cases (min = max)
+gotests -all -ai -ai-min-cases 5 -ai-max-cases 5 -w yourfile.go
+
+# Generate a range of test cases (AI chooses between 3-7)
+gotests -all -ai -ai-min-cases 3 -ai-max-cases 7 -w yourfile.go
+
+# Combine with other flags
+gotests -exported -ai -parallel -w yourfile.go
+```
+
+### How It Works
+
+- Analyzes function implementation and logic
+- Generates realistic test values based on actual code
+- Creates test cases for edge cases and error conditions
+- Falls back to TODO comments if generation fails
+- Works offline with local models (privacy-first)
+
+### Supported Features
+
+✅ Simple types (int, string, bool, float)
+✅ Complex types (slices, maps, structs, pointers)
+✅ Error returns and validation
+✅ Variadic parameters
+✅ Methods with receivers
+✅ Multiple return values
+
+### Privacy & Security
+
+**What data is sent to the LLM:**
+- Function signatures (name, parameters, return types)
+- Complete function bodies including all code and comments
+- No file paths or project context
+
+**Privacy considerations:**
+- ⚠️ **Function bodies may contain sensitive information** - business logic, algorithms, or credentials/secrets in comments
+- ✅ **Local-first by default** - Using Ollama keeps all data on your machine; nothing is sent to external servers
+- ✅ **Offline operation** - AI generation works completely offline with local models
+- 🔒 **Recommendation**: Avoid using `-ai` on code containing secrets, API keys, or proprietary algorithms in comments
+
+**If using cloud providers in the future:**
+- Function source code will be transmitted to the cloud provider's API
+- Review the provider's data retention and privacy policies
+- Consider using `-ai` only on non-sensitive codebases
+
 ## Quick Start Examples
 
 ### Generate tests for a single function

gotests.go

@@ -31,6 +31,11 @@ type Options struct {
 	TemplateParams map[string]interface{} // Custom external parameters
 	TemplateData   [][]byte               // Data slice for templates
 	UseGoCmp       bool                   // Use cmp.Equal (google/go-cmp) instead of reflect.DeepEqual
+	UseAI          bool                   // Generate test cases using AI
+	AIModel        string                 // AI model to use
+	AIEndpoint     string                 // AI API endpoint
+	AIMinCases     int                    // Minimum number of test cases to generate
+	AIMaxCases     int                    // Maximum number of test cases to generate
 }
 
 // A GeneratedTest contains information about a test file with generated tests.
@@ -131,6 +136,11 @@ func generateTest(src models.Path, files []models.Path, opt *Options) (*Generate
 		TemplateDir:    opt.TemplateDir,
 		TemplateParams: opt.TemplateParams,
 		TemplateData:   opt.TemplateData,
+		UseAI:          opt.UseAI,
+		AIModel:        opt.AIModel,
+		AIEndpoint:     opt.AIEndpoint,
+		AIMinCases:     opt.AIMinCases,
+		AIMaxCases:     opt.AIMaxCases,
 	}
 
 	b, err := options.Process(h, funcs)

gotests/main.go

@@ -63,6 +63,11 @@ var (
 	templateParamsPath = flag.String("template_params_file", "", "read external parameters to template by json with file")
 	templateParams     = flag.String("template_params", "", "read external parameters to template by json with stdin")
 	useGoCmp           = flag.Bool("use_go_cmp", false, "use cmp.Equal (google/go-cmp) instead of reflect.DeepEqual")
+	useAI              = flag.Bool("ai", false, "generate test cases using AI (requires Ollama)")
+	aiModel            = flag.String("ai-model", "qwen2.5-coder:0.5b", "AI model to use for test generation")
+	aiEndpoint         = flag.String("ai-endpoint", "http://localhost:11434", "Ollama API endpoint")
+	aiMinCases         = flag.Int("ai-min-cases", 3, "minimum number of test cases to generate with AI")
+	aiMaxCases         = flag.Int("ai-max-cases", 10, "maximum number of test cases to generate with AI")
 	version            = flag.Bool("version", false, "print version information and exit")
 )
 
@@ -88,6 +93,31 @@ func main() {
 		return
 	}
 
+	// Validate AI parameters and warn user
+	if *useAI {
+		// Warn about sending code to AI provider
+		fmt.Fprintf(os.Stderr, "⚠️  WARNING: Function source code will be sent to AI provider at %s\n", *aiEndpoint)
+		fmt.Fprintf(os.Stderr, "   Ensure your code does not contain secrets or sensitive information.\n\n")
+
+		// Validate parameters
+		if *aiModel == "" {
+			fmt.Fprintf(os.Stderr, "Error: -ai-model cannot be empty when using -ai flag\n")
+			os.Exit(1)
+		}
+		if *aiMinCases < 1 {
+			fmt.Fprintf(os.Stderr, "Error: -ai-min-cases must be at least 1, got %d\n", *aiMinCases)
+			os.Exit(1)
+		}
+		if *aiMaxCases > 100 {
+			fmt.Fprintf(os.Stderr, "Error: -ai-max-cases must be at most 100, got %d\n", *aiMaxCases)
+			os.Exit(1)
+		}
+		if *aiMinCases > *aiMaxCases {
+			fmt.Fprintf(os.Stderr, "Error: -ai-min-cases (%d) cannot be greater than -ai-max-cases (%d)\n", *aiMinCases, *aiMaxCases)
+			os.Exit(1)
+		}
+	}
+
 	process.Run(os.Stdout, args, &process.Options{
 		OnlyFuncs:          *onlyFuncs,
 		ExclFuncs:          *exclFuncs,
@@ -103,6 +133,11 @@ func main() {
 		TemplateParamsPath: *templateParamsPath,
 		TemplateParams:     *templateParams,
 		UseGoCmp:           *useGoCmp,
+		UseAI:              *useAI,
+		AIModel:            *aiModel,
+		AIEndpoint:         *aiEndpoint,
+		AIMinCases:         *aiMinCases,
+		AIMaxCases:         *aiMaxCases,
 	})
 }
 

gotests/process/process.go

@@ -38,6 +38,11 @@ type Options struct {
 	TemplateParams     string   // Custom parameters as JSON string
 	TemplateData       [][]byte // Data slice for templates
 	UseGoCmp           bool     // Use cmp.Equal (google/go-cmp) instead of reflect.DeepEqual
+	UseAI              bool     // Generate test cases using AI
+	AIModel            string   // AI model to use
+	AIEndpoint         string   // AI API endpoint
+	AIMinCases         int      // Minimum number of test cases to generate
+	AIMaxCases         int      // Maximum number of test cases to generate
 }
 
 // Run generates tests for the Go files defined in args with the given options.
@@ -116,6 +121,11 @@ func parseOptions(out io.Writer, opt *Options) *gotests.Options {
 		TemplateParams: templateParams,
 		TemplateData:   opt.TemplateData,
 		UseGoCmp:       opt.UseGoCmp,
+		UseAI:          opt.UseAI,
+		AIModel:        opt.AIModel,
+		AIEndpoint:     opt.AIEndpoint,
+		AIMinCases:     opt.AIMinCases,
+		AIMaxCases:     opt.AIMaxCases,
 	}
 }