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;
        private string _sceneObjectLogicalName;
        private string _prefabLogicalName;

        [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);
            
            Directory.CreateDirectory(Path.GetDirectoryName(TestPrefabPath));
            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();
            _sceneObjectLogicalName = "scene_object_1";
            _prefabLogicalName = "prefab_asset_1";

            var scenePid = new PersistentIdentifier
            {
                logicalName = _sceneObjectLogicalName,
                instanceID = _testSceneObject.GetInstanceID(),
                path = CommandUtility.GetGameObjectPath(_testSceneObject)
            };
            _context.IdentifierMap[_sceneObjectLogicalName] = scenePid;

            var prefabGuid = AssetDatabase.AssetPathToGUID(TestPrefabPath);
            var prefabPid = new PersistentIdentifier
            {
                logicalName = _prefabLogicalName,
                guid = prefabGuid
            };
            _context.IdentifierMap[_prefabLogicalName] = prefabPid;
        }

        [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(_sceneObjectLogicalName, "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(_sceneObjectLogicalName, "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 command = CreateCommand(_prefabLogicalName, "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 command = CreateCommand(_prefabLogicalName, "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(_sceneObjectLogicalName, "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"));
        }
        
        [Test]
        public void Execute_AddRigidbodyToSceneObject_Successfully()
        {
            // Arrange
            var command = CreateCommand(_sceneObjectLogicalName, "UnityEngine.Rigidbody");

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

            // Assert
            Assert.AreEqual(CommandOutcome.Success, outcome);
            Assert.IsNotNull(_testSceneObject.GetComponent<Rigidbody>());
        }
        
        [Test]
        public void Execute_AddRigidbodyToSceneObjectWithExistingRigidbody_ReturnsError()
        {
            // Arrange
            _testSceneObject.AddComponent<Rigidbody>();
            var command = CreateCommand(_sceneObjectLogicalName, "UnityEngine.Rigidbody");

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

            // Assert
            Assert.AreEqual(CommandOutcome.Error, outcome);
            Assert.IsNotNull(_context.ErrorMessage);
            Assert.IsTrue(_context.ErrorMessage.Contains("Failed to add component 'UnityEngine.Rigidbody'"));
        }
    }
}