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

using NUnit.Framework;
using UnityEngine;
using UnityEditor;
using System.IO;
using LLM.Editor.Commands;
using LLM.Editor.Data;
using LLM.Editor.Helper;

namespace LLM.Editor.Tests.Unit
{
    public class MyTestCustomComponent : MonoBehaviour {}

    [TestFixture]
    public class AddComponentToAssetCommandTests
    {
        private const string TestPrefabPath = "Assets/LLM/Editor/Tests/Prefabs/TempTestPrefab.prefab";
        private GameObject _testSceneObject;
        private CommandContext _context;

        [SetUp]
        public void SetUp()
        {
            // Create a scene object for scene-based tests, with a child
            _testSceneObject = new GameObject("TestSceneObject");
            var sceneChild = new GameObject("SceneChild");
            sceneChild.transform.SetParent(_testSceneObject.transform);
            
            // Create a prefab asset for prefab-based tests
            var prefabSourceObject = new GameObject("TestPrefab");
            var prefabChild = new GameObject("PrefabChild");
            prefabChild.transform.SetParent(prefabSourceObject.transform);
            PrefabUtility.SaveAsPrefabAsset(prefabSourceObject, TestPrefabPath);
            Object.DestroyImmediate(prefabSourceObject);

            // Set up the context with identifiers for both the scene object and the prefab asset
            _context = new CommandContext
            {
                IdentifierMap =
                {
                    [_testSceneObject.GetInstanceID().ToString()] = _testSceneObject.GetInstanceID().ToString(),
                    [AssetDatabase.AssetPathToGUID(TestPrefabPath)] = AssetDatabase.AssetPathToGUID(TestPrefabPath)
                }
            };
        }

        [TearDown]
        public void TearDown()
        {
            if (_testSceneObject != null) Object.DestroyImmediate(_testSceneObject);
            AssetDatabase.DeleteAsset(TestPrefabPath);
        }

        private static AddComponentToAssetCommand CreateCommand(string targetId, string scriptName, string childPath = null)
        {
            var parameters = new AddComponentParams
            {
                targetIdentifier = targetId,
                scriptName = scriptName,
                childPath = childPath
            };
            return new AddComponentToAssetCommand(parameters.ToJson());
        }

        [Test]
        public void Execute_OnSceneObjectRoot_AddsComponentSuccessfully()
        {
            // Arrange
            var command = CreateCommand(_testSceneObject.GetInstanceID().ToString(), "UnityEngine.Rigidbody");

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

            // Assert
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.IsNotNull(_testSceneObject.GetComponent<Rigidbody>());
        }

        [Test]
        public void Execute_OnSceneObjectWithValidChildPath_AddsComponentToChild()
        {
            // Arrange
            var command = CreateCommand(_testSceneObject.GetInstanceID().ToString(), "UnityEngine.BoxCollider", "SceneChild");

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

            // Assert
            var childObject = _testSceneObject.transform.Find("SceneChild");
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.IsNotNull(childObject.GetComponent<BoxCollider>());
            Assert.IsNull(_testSceneObject.GetComponent<BoxCollider>()); // Ensure it wasn't added to the parent
        }

        [Test]
        public void Execute_OnPrefabRoot_AddsComponentToRootInAsset()
        {
            // Arrange
            var prefabGuid = AssetDatabase.AssetPathToGUID(TestPrefabPath);
            var command = CreateCommand(prefabGuid, "UnityEngine.SphereCollider");

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

            // Assert
            Assert.AreEqual(CommandOutcome.Success, outcome);
            var loadedPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(TestPrefabPath);
            Assert.IsNotNull(loadedPrefab.GetComponent<SphereCollider>(), "Component should be added to the prefab asset's root.");
        }

        [Test]
        public void Execute_OnPrefabWithValidChildPath_AddsComponentToChildInAsset()
        {
            // Arrange
            var prefabGuid = AssetDatabase.AssetPathToGUID(TestPrefabPath);
            var command = CreateCommand(prefabGuid, "UnityEngine.CapsuleCollider", "PrefabChild");

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

            // Assert
            Assert.AreEqual(CommandOutcome.Success, outcome);
            var loadedPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(TestPrefabPath);
            var childInPrefab = loadedPrefab.transform.Find("PrefabChild");
            Assert.IsNotNull(childInPrefab, "Child object should exist in prefab.");
            Assert.IsNotNull(childInPrefab.GetComponent<CapsuleCollider>(), "Component should be added to the child within the prefab asset.");
        }

        [Test]
        public void Execute_WithInvalidScriptName_ReturnsError()
        {
            // Arrange
            var command = CreateCommand(_testSceneObject.GetInstanceID().ToString(), "Invalid.NonExistent.Component");

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

            // Assert
            Assert.AreEqual(CommandOutcome.Error, outcome);
            Assert.IsNotNull(_context.ErrorMessage);
            Assert.IsTrue(_context.ErrorMessage.Contains("Could not find script type"));
        }
    }
}