From b6c8060af12ca9fb88bfadfd8b25a3101ebb4bfc Mon Sep 17 00:00:00 2001
From: Leon Friedrich <60421075+ElectroJr@users.noreply.github.com>
Date: Sun, 27 Aug 2023 00:23:32 +1200
Subject: [PATCH] Add new PVS test (#4312)

---
 .../ComponentNetworkGenerator.cs              |   2 +-
 Robust.UnitTesting/Robust.UnitTesting.csproj  |   1 +
 .../Shared/GameState/ComponentStateTests.cs   | 153 ++++++++++++++++++
 3 files changed, 155 insertions(+), 1 deletion(-)
 create mode 100644 Robust.UnitTesting/Shared/GameState/ComponentStateTests.cs

diff --git a/Robust.Shared.CompNetworkGenerator/ComponentNetworkGenerator.cs b/Robust.Shared.CompNetworkGenerator/ComponentNetworkGenerator.cs
index a486ca8e71..20e1da16a3 100644
--- a/Robust.Shared.CompNetworkGenerator/ComponentNetworkGenerator.cs
+++ b/Robust.Shared.CompNetworkGenerator/ComponentNetworkGenerator.cs
@@ -158,7 +158,7 @@ namespace {nameSpace};
 
 public partial class {componentName}
 {{
-    [Serializable, NetSerializable]
+    [System.Serializable, NetSerializable]
     public class {stateName} : ComponentState
     {{{stateFields}
     }}
diff --git a/Robust.UnitTesting/Robust.UnitTesting.csproj b/Robust.UnitTesting/Robust.UnitTesting.csproj
index 64d9def3d9..1f51dee7b6 100644
--- a/Robust.UnitTesting/Robust.UnitTesting.csproj
+++ b/Robust.UnitTesting/Robust.UnitTesting.csproj
@@ -33,4 +33,5 @@
     </EmbeddedResource>
   </ItemGroup>
   <Import Project="..\MSBuild\Robust.Properties.targets" />
+  <Import Project="..\MSBuild\Robust.CompNetworkGenerator.targets" />
 </Project>
diff --git a/Robust.UnitTesting/Shared/GameState/ComponentStateTests.cs b/Robust.UnitTesting/Shared/GameState/ComponentStateTests.cs
new file mode 100644
index 0000000000..ce115c9797
--- /dev/null
+++ b/Robust.UnitTesting/Shared/GameState/ComponentStateTests.cs
@@ -0,0 +1,153 @@
+using System.Linq;
+using System.Numerics;
+using System.Threading.Tasks;
+using NUnit.Framework;
+using Robust.Server.GameObjects;
+using Robust.Server.Player;
+using Robust.Shared;
+using Robust.Shared.GameObjects;
+using Robust.Shared.GameStates;
+using Robust.Shared.IoC;
+using Robust.Shared.Map;
+using Robust.Shared.Network;
+
+namespace Robust.UnitTesting.Shared.GameState;
+
+public sealed partial class ComponentStateTests : RobustIntegrationTest
+{
+    /// <summary>
+    /// This tests performs a basic check to ensure that there is no issue with entity states referencing other
+    /// entities that the client is not yet aware of. It does this by spawning two entities that reference each other,
+    /// and then ensuring that they get sent to the client one at a time.
+    /// </summary>
+    [Test]
+    public async Task UnknownEntityTest()
+    {
+        // Setup auto-comp-states. I hate this. Someone please fix reflection in RobustIntegrationTest
+        var compReg = () => IoCManager.Resolve<IComponentFactory>().RegisterClass<UnknownEntityTestComponent>();
+        var sysReg = () => IoCManager.Resolve<IEntitySystemManager>().LoadExtraSystemType<UnknownEntityTestComponent.UnknownEntityTestComponent_AutoNetworkSystem>();
+        var serverOpts = new ServerIntegrationOptions
+        {
+            Pool = false,
+            BeforeRegisterComponents = compReg,
+            BeforeStart = sysReg,
+        };
+        var clientOpts = new ClientIntegrationOptions
+        {
+            Pool = false,
+            BeforeRegisterComponents = compReg,
+            BeforeStart = sysReg,
+        };
+        var server = StartServer(serverOpts);
+        var client = StartClient(clientOpts);
+
+        await Task.WhenAll(client.WaitIdleAsync(), server.WaitIdleAsync());
+        var netMan = client.ResolveDependency<IClientNetManager>();
+        var xforms = server.System<SharedTransformSystem>();
+
+        Assert.DoesNotThrow(() => client.SetConnectTarget(server));
+        client.Post(() => netMan.ClientConnect(null!, 0, null!));
+        server.Post(() => server.CfgMan.SetCVar(CVars.NetPVS, true));
+
+        // Set up map.
+        EntityUid map = default;
+        await server.WaitPost(() =>
+        {
+            var mapId = server.MapMan.CreateMap();
+            map = server.MapMan.GetMapEntityId(mapId);
+        });
+
+        await RunTicks();
+
+        // Spawn entities
+        var coordsA = new EntityCoordinates(map, default);
+        var coordsB = new EntityCoordinates(map, new Vector2(100, 100));
+        EntityUid player = default;
+        EntityUid entA = default;
+        EntityUid entB = default;
+
+        await server.WaitPost(() =>
+        {
+            // Attach player.
+            player = server.EntMan.Spawn();
+            var session = (IPlayerSession) server.PlayerMan.Sessions.First();
+            server.System<ActorSystem>().Attach(player, session);
+            session.JoinGame();
+
+            // Spawn test entities.
+            entA = server.EntMan.SpawnAttachedTo(null, coordsA);
+            entB = server.EntMan.SpawnAttachedTo(null, coordsB);
+
+            // Setup components
+            var cmp = server.EntMan.EnsureComponent<UnknownEntityTestComponent>(entA);
+            cmp.Other = entB;
+            server.EntMan.Dirty(entA, cmp);
+
+            cmp = server.EntMan.EnsureComponent<UnknownEntityTestComponent>(entB);
+            cmp.Other = entA;
+            server.EntMan.Dirty(entB, cmp);
+        });
+
+        await RunTicks();
+
+        // Check player got properly attached and only knows about the expected entities
+        await client.WaitPost(() =>
+        {
+            Assert.That(client.AttachedEntity, Is.EqualTo(player));
+            Assert.That(client.EntMan.EntityExists(player));
+            Assert.That(client.EntMan.EntityExists(entA), Is.False);
+            Assert.That(client.EntMan.EntityExists(entB), Is.False);
+        });
+
+        // Move the player into PVS range of one of the entities.
+        await server.WaitPost(() => xforms.SetCoordinates(player, coordsB));
+        await RunTicks();
+
+        await client.WaitPost(() =>
+        {
+            Assert.That(client.EntMan.EntityExists(entB), Is.True);
+            Assert.That(client.EntMan.EntityExists(entA), Is.False);
+
+            Assert.That(client.EntMan.TryGetComponent(entB, out UnknownEntityTestComponent? cmp));
+            Assert.That(cmp?.Other, Is.EqualTo(entA));
+        });
+
+        // Move the player into PVS range of the other entity
+        await server.WaitPost(() => xforms.SetCoordinates(player, coordsA));
+        await RunTicks();
+
+        await client.WaitPost(() =>
+        {
+            Assert.That(client.EntMan.EntityExists(entB), Is.True);
+            Assert.That(client.EntMan.EntityExists(entA), Is.True);
+
+            Assert.That(client.EntMan.TryGetComponent(entB, out UnknownEntityTestComponent? cmp));
+            Assert.That(cmp?.Other, Is.EqualTo(entA));
+
+            Assert.That(client.EntMan.TryGetComponent(entA, out cmp));
+            Assert.That(cmp?.Other, Is.EqualTo(entB));
+        });
+
+        server.Post(() => server.CfgMan.SetCVar(CVars.NetPVS, false));
+
+        // wait for errors.
+        await RunTicks();
+
+        async Task RunTicks()
+        {
+            for (int i = 0; i < 10; i++)
+            {
+                await server!.WaitRunTicks(1);
+                await client!.WaitRunTicks(1);
+            }
+        }
+    }
+
+}
+
+[RegisterComponent, NetworkedComponent, AutoGenerateComponentState]
+public sealed partial class UnknownEntityTestComponent : Component
+{
+    [AutoNetworkedField]
+    public EntityUid? Other;
+}
\ No newline at end of file