feat(agent): Support updating existing unit tests (#31)

README.md

@@ -108,6 +108,20 @@ Options:
 
 This command will generate a unit test for the specified code file, optionally using an example unit test file as a reference.
 
+#### Updating Unit Tests
+
+To update an existing unit test with code changes, use the `unittest update` command:
+
+```bash
+flowpair unittest update [options]
+```
+
+Options:
+- `-s` or `--source-file`: The file path of the code to test (Required)
+- `-t` or `--test-file`: The file path of the existing unit tests (Required)
+
+This command will update the existing unit test file to reflect changes made in the source code file.
+
 ## Contributing
 
 We welcome contributions to FlowPair! If you have suggestions for improvements or encounter any issues, please feel free to:

src/FlowPair/Agent/Infrastructure/IAgentModule.cs

@@ -2,6 +2,7 @@
 using Raiqub.LlmTools.FlowPair.Agent.Operations.CreateUnitTest;
 using Raiqub.LlmTools.FlowPair.Agent.Operations.Login;
 using Raiqub.LlmTools.FlowPair.Agent.Operations.ReviewChanges;
+using Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest;
 using Raiqub.LlmTools.FlowPair.Agent.Services;
 
 namespace Raiqub.LlmTools.FlowPair.Agent.Infrastructure;
@@ -18,12 +19,14 @@ namespace Raiqub.LlmTools.FlowPair.Agent.Infrastructure;
 // Chat definitions
 [Singleton(typeof(IReviewChatDefinition), typeof(ReviewChatDefinition))]
 [Singleton(typeof(ICreateUnitTestChatDefinition), typeof(CreateUnitTestChatDefinition))]
+[Singleton(typeof(IUpdateUnitTestChatDefinition), typeof(UpdateUnitTestChatDefinition))]
 
 // Operations
 [Singleton(typeof(ILoginUseCase), typeof(LoginUseCase))]
 [Singleton(typeof(LoginCommand))]
 [Singleton(typeof(ReviewChangesCommand))]
 [Singleton(typeof(CreateUnitTestCommand))]
+[Singleton(typeof(UpdateUnitTestCommand))]
 public interface IAgentModule
 {
     static AgentJsonContext GetJsonContext() => AgentJsonContext.Default;

src/FlowPair/Agent/Operations/UpdateUnitTest/UpdateUnitTestChatDefinition.cs

@@ -0,0 +1,70 @@
+using Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest.v1;
+using Raiqub.LlmTools.FlowPair.Chats.Contracts.v1;
+using Raiqub.LlmTools.FlowPair.Chats.Models;
+using Raiqub.LlmTools.FlowPair.Chats.Services;
+
+namespace Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest;
+
+public interface IUpdateUnitTestChatDefinition : IChatDefinition<UpdateUnitTestResponse>;
+
+public sealed class UpdateUnitTestChatDefinition
+    : IUpdateUnitTestChatDefinition
+{
+    private const string CodeResponseKey = "Markdown";
+
+    public ChatScript ChatScript { get; } = new(
+        "Update unit tests chat script",
+        [
+            /* Python          */".py", ".pyw", ".pyx", ".pxd", ".pxi",
+            /* JavaScript      */".js", ".jsx", ".mjs", ".cjs",
+            /* Java            */".java",
+            /* C#              */".cs", ".csx",
+            /* C++             */".cpp", ".cxx", ".cc", ".c++", ".hpp", ".hxx", ".h", ".hh", ".h++",
+            /* PHP             */".php", ".phtml", ".phps",
+            /* Ruby            */".rb", ".rbw", ".rake",
+            /* Swift           */".swift",
+            /* R               */".r",
+            /* SQL             */".sql",
+            /* Kotlin          */".kt", ".kts",
+            /* TypeScript      */".ts", ".tsx",
+            /* Go (Golang)     */".go",
+            /* Rust            */".rs",
+            /* Scala           */".scala", ".sc",
+            /* Dart            */".dart",
+            /* Perl            */".pl", ".pm", ".t", ".pod",
+            /* MATLAB          */".m",
+        ],
+        """
+        You are an expert developer, your task is to create unit tests following the best practices.
+        You are given a set of project files, containing the filenames and their contents.
+        """,
+        [
+            Instruction.StepInstruction.Of(
+                "Update the unit tests for the specified source code"),
+            Instruction.StepInstruction.Of(
+                "Ensure the unit tests cover every possible execution path in the code"),
+            Instruction.StepInstruction.Of(
+                "Ensure the unit tests are sensitive to mutations in the source code. " +
+                "When mutation testing introduces small changes to the implementation (mutants), " +
+                "at least one test should fail. " +
+                "This verifies that the tests can detect potential bugs or behavioral changes."),
+            Instruction.StepInstruction.Of(
+                "Remove any redundant tests while maintaining full coverage"),
+            Instruction.CodeExtractInstruction.Of(
+                CodeResponseKey,
+                "Return the entire final version of the updated unit tests file content, " +
+                "incorporating all the above improvements, inside a code block (```)"),
+        ]);
+
+    public Result<object, string> Parse(string key, string input) => key switch
+    {
+        CodeResponseKey => MarkdownCodeExtractor
+            .TryExtractSingle(input)
+            .Select(static object (x) => x),
+        _ => $"Unknown output key '{key}'"
+    };
+
+    public Option<UpdateUnitTestResponse> ConvertResult(ChatWorkspace chatWorkspace) =>
+        from code in OutputProcessor.GetFirst<CodeSnippet>(chatWorkspace, CodeResponseKey)
+        select new UpdateUnitTestResponse(code.Content);
+}

src/FlowPair/Agent/Operations/UpdateUnitTest/UpdateUnitTestCommand.cs

@@ -0,0 +1,109 @@
+using System.IO.Abstractions;
+using System.Text;
+using ConsoleAppFramework;
+using Raiqub.LlmTools.FlowPair.Agent.Operations.Login;
+using Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest.v1;
+using Raiqub.LlmTools.FlowPair.Agent.Services;
+using Raiqub.LlmTools.FlowPair.Chats.Models;
+using Raiqub.LlmTools.FlowPair.Chats.Services;
+using Raiqub.LlmTools.FlowPair.Common;
+using Raiqub.LlmTools.FlowPair.LocalFileSystem.Services;
+using Spectre.Console;
+
+namespace Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest;
+
+public sealed class UpdateUnitTestCommand(
+    IAnsiConsole console,
+    IFileSystem fileSystem,
+    IUpdateUnitTestChatDefinition chatDefinition,
+    IWorkingDirectoryWalker workingDirectoryWalker,
+    IProjectFilesMessageFactory projectFilesMessageFactory,
+    IDirectoryStructureMessageFactory directoryStructureMessageFactory,
+    ILoginUseCase loginUseCase,
+    IChatService chatService)
+{
+    /// <summary>
+    /// Update existing unit test with code changes.
+    /// </summary>
+    /// <param name="sourceFile">-s, The file path of the code to test.</param>
+    /// <param name="testFile">-t, The file path of the existing unit tests.</param>
+    [Command("unittest update")]
+    public int Execute(
+        string sourceFile,
+        string testFile)
+    {
+        var sourceFileInfo = fileSystem.FileInfo.New(PathAnalyzer.Normalize(sourceFile));
+        if (!sourceFileInfo.Exists)
+        {
+            console.MarkupLine("[red]Error:[/] The specified source file does not exist.");
+            return 1;
+        }
+
+        var testFileInfo = fileSystem.FileInfo.New(PathAnalyzer.Normalize(testFile));
+        if (!testFileInfo.Exists)
+        {
+            console.MarkupLine("[red]Error:[/] The specified test file does not exist.");
+            return 2;
+        }
+
+        return (from rootPath in workingDirectoryWalker.TryFindRepositoryRoot(sourceFileInfo.Directory?.FullName)
+                    .OkOrElse(HandleFindRepositoryRootError)
+                from session in loginUseCase.Execute(isBackground: true)
+                    .UnwrapErrOr(0)
+                    .Ensure(n => n == 0, 4)
+                let initialMessages = BuildInitialMessages(sourceFileInfo, testFileInfo, rootPath).ToList()
+                from response in chatService
+                    .Run(console.Progress(), LlmModelType.Claude35Sonnet, chatDefinition, initialMessages)
+                    .MapErr(HandleChatServiceError)
+                let create = RewriteUnitTestFile(rootPath, testFileInfo, response)
+                select 0)
+            .UnwrapEither();
+    }
+
+    private int HandleFindRepositoryRootError()
+    {
+        console.MarkupLine("[red]Error:[/] Could not locate Git repository.");
+        return 3;
+    }
+
+    private int HandleChatServiceError(string errorMessage)
+    {
+        console.MarkupLineInterpolated($"[red]Error:[/] {errorMessage}");
+        return 5;
+    }
+
+    private IEnumerable<Message> BuildInitialMessages(
+        IFileInfo sourceFileInfo,
+        IFileInfo testFileInfo,
+        IDirectoryInfo rootPath)
+    {
+        yield return projectFilesMessageFactory.CreateWithProjectFilesContent(rootPath);
+        yield return directoryStructureMessageFactory.CreateWithRepositoryStructure(rootPath);
+
+        yield return new Message(
+            SenderRole.User,
+            $"""
+             The source file updated content is:
+             ```
+             {sourceFileInfo.ReadAllText()}
+             ```
+             """);
+
+        yield return new Message(
+            SenderRole.User,
+            $"""
+             The existing test file content is:
+             ```
+             {testFileInfo.ReadAllText()}
+             ```
+             """);
+    }
+
+    private Unit RewriteUnitTestFile(IDirectoryInfo rootPath, IFileInfo testFileInfo, UpdateUnitTestResponse response)
+    {
+        var fileRelativePath = rootPath.GetRelativePath(testFileInfo.FullName);
+        testFileInfo.WriteAllText(response.Content, Encoding.UTF8);
+        console.MarkupLineInterpolated($"[green]Unit tests updated:[/] {fileRelativePath}");
+        return Unit();
+    }
+}

src/FlowPair/Agent/Operations/UpdateUnitTest/v1/UpdateUnitTestResponse.cs

@@ -0,0 +1,4 @@
+namespace Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest.v1;
+
+public sealed record UpdateUnitTestResponse(
+    string Content);

src/FlowPair/FlowPair.csproj

@@ -38,6 +38,10 @@
     <PackageReference Include="TestableIO.System.IO.Abstractions.Wrappers" />
   </ItemGroup>
 
+  <ItemGroup>
+    <InternalsVisibleTo Include="Raiqub.LlmTools.FlowPair.Tests" />
+  </ItemGroup>
+
   <ItemGroup>
     <Compile Update="Agent\Operations\ReviewChanges\FeedbackHtmlTemplate.cs">
       <AutoGen>True</AutoGen>

src/FlowPair/Program.cs

@@ -2,16 +2,38 @@
 using Raiqub.LlmTools.FlowPair.Agent.Operations.CreateUnitTest;
 using Raiqub.LlmTools.FlowPair.Agent.Operations.Login;
 using Raiqub.LlmTools.FlowPair.Agent.Operations.ReviewChanges;
+using Raiqub.LlmTools.FlowPair.Agent.Operations.UpdateUnitTest;
 using Raiqub.LlmTools.FlowPair.DependencyInjection;
 using Raiqub.LlmTools.FlowPair.Settings.Operations.Configure;
+using Spectre.Console;
 
-using var container = new AppContainer();
-ConsoleApp.ServiceProvider = container;
-ConsoleApp.Version = ThisAssembly.AssemblyInformationalVersion;
+namespace Raiqub.LlmTools.FlowPair;
 
-var app = ConsoleApp.Create();
-app.Add<ConfigureCommand>();
-app.Add<LoginCommand>();
-app.Add<ReviewChangesCommand>();
-app.Add<CreateUnitTestCommand>();
-app.Run(args);
+public static class Program
+{
+    private static readonly Style s_errorStyle = new(Color.Red);
+
+    public static void Main(string[] args)
+    {
+        using var container = new AppContainer();
+        var console = container.GetService<IAnsiConsole>();
+
+        ConsoleApp.ServiceProvider = container;
+        ConsoleApp.Version = ThisAssembly.AssemblyInformationalVersion;
+        ConsoleApp.Log = x => console.WriteLine(x);
+        ConsoleApp.LogError = x => console.WriteLine(x, s_errorStyle);
+
+        RunConsoleApp(args);
+    }
+
+    public static void RunConsoleApp(string[] args)
+    {
+        var app = ConsoleApp.Create();
+        app.Add<ConfigureCommand>();
+        app.Add<LoginCommand>();
+        app.Add<ReviewChangesCommand>();
+        app.Add<CreateUnitTestCommand>();
+        app.Add<UpdateUnitTestCommand>();
+        app.Run(args);
+    }
+}