terra-llm/Assets/LLM/Editor/Tests/Unit/UpdateWorkingContextCommandTests.cs

using NUnit.Framework;
using LLM.Editor.Commands;
using LLM.Editor.Data;
using LLM.Editor.Core;
using Newtonsoft.Json.Linq;
using System.Collections.Generic;

namespace LLM.Editor.Tests.Unit
{
    [TestFixture]
    public class UpdateWorkingContextCommandTests
    {
        private CommandContext _context;

        [SetUp]
        public void SetUp()
        {
            _context = new CommandContext();
            // Start a session to ensure file operations in SessionManager work correctly.
            SessionManager.StartNewSession();
            // Ensure a clean slate for each test
            SessionManager.SaveWorkingContext(new JObject());
        }

        [TearDown]
        public void TearDown()
        {
            // Clean up session state after tests
            SessionManager.EndSession();
        }

        [Test]
        public void Execute_WithUpdates_MergesDataIntoWorkingContext()
        {
            // Arrange
            var initialContext = new JObject { ["existingKey"] = "existingValue" };
            SessionManager.SaveWorkingContext(initialContext);

            var jsonParams = @"{
                ""updates"": {
                    ""newKey"": ""newValue"",
                    ""existingKey"": ""updatedValue""
                }
            }";
            var command = new UpdateWorkingContextCommand(jsonParams);

            // Act
            var outcome = command.Execute(_context);

            // Assert
            var finalContext = SessionManager.LoadWorkingContext();
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.AreEqual("newValue", (string)finalContext["newKey"]);
            Assert.AreEqual("updatedValue", (string)finalContext["existingKey"]);
        }

        [Test]
        public void Execute_WithPlanUpdate_PopulatesRuntimeContextPlan()
        {
            // Arrange
            var jsonParams = @"{
                ""updates"": {
                    ""plan"": [""Step 1"", ""Step 2""]
                }
            }";
            var command = new UpdateWorkingContextCommand(jsonParams);

            // Act
            var outcome = command.Execute(_context);

            // Assert
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.IsNotNull(_context.Plan);
            Assert.AreEqual(2, _context.Plan.Count);
            Assert.AreEqual("Step 1", _context.Plan[0]);
        }

        [Test]
        public void Execute_WithKeysToRemove_RemovesDataFromWorkingContext()
        {
            // Arrange
            var initialContext = new JObject { ["key1"] = "value1", ["key2"] = "value2" };
            SessionManager.SaveWorkingContext(initialContext);

            const string jsonParams = @"{
                ""keysToRemove"": [""key1""]
            }";
            var command = new UpdateWorkingContextCommand(jsonParams);

            // Act
            var outcome = command.Execute(_context);

            // Assert
            var finalContext = SessionManager.LoadWorkingContext();
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.IsFalse(finalContext.ContainsKey("key1"));
            Assert.IsTrue(finalContext.ContainsKey("key2"));
        }

        [Test]
        public void Execute_WithClearAll_WipesWorkingContext()
        {
            // Arrange
            var initialContext = new JObject { ["key1"] = "value1" };
            SessionManager.SaveWorkingContext(initialContext);
            _context.Plan = new List<string> { "an old plan" };

            var jsonParams = @"{ ""clearAll"": true }";
            var command = new UpdateWorkingContextCommand(jsonParams);

            // Act
            var outcome = command.Execute(_context);

            // Assert
            var finalContext = SessionManager.LoadWorkingContext();
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.AreEqual(0, finalContext.Count);
            Assert.IsNull(_context.Plan, "Runtime context plan should also be cleared.");
        }
    }
}