update TaskingTests to be more reliable

MythicAgents · Feb 16, 2024 · 98aa74a · 98aa74a
1 parent c1522c3
commit 98aa74a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/AgentTests/TaskingTests.cs b/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/AgentTests/TaskingTests.cs
@@ -30,9 +30,11 @@ public async Task TestGetTaskingSingle()
             TestProfile prof = (TestProfile)_profile.First();
 
             Task.Run(() => _agent.Start());
-            prof.taskingSent.WaitOne(1000);
+            ((TestTaskManager)_taskManager).WaitForNumberOfJobs(1);
+            //prof.taskingSent.WaitOne(1000);
             _profile.First().StopBeacon();
-            Assert.IsTrue(((TestTaskManager)_taskManager).jobs.Count > 0);
+            Console.WriteLine(((TestTaskManager)_taskManager).jobs.Count);
+            Assert.IsTrue(((TestTaskManager)_taskManager).jobs.Count == 1);
         }
         [TestMethod]
         public async Task TestGetTaskingMultiple()
@@ -79,9 +81,11 @@ public async Task TestGetTaskingMultiple()
             TestProfile prof = (TestProfile)_profile.First();
 
             Task.Run(_agent.Start);
-            prof.taskingSent.WaitOne(50);
+            ((TestTaskManager)_taskManager).WaitForNumberOfJobs(3);
+            //prof.taskingSent.WaitOne(1000);
             _profile.First().StopBeacon();
-            Assert.IsTrue(((TestTaskManager)_taskManager).jobs.Count > 1);
+            Console.WriteLine(((TestTaskManager)_taskManager).jobs.Count);
+            Assert.IsTrue(((TestTaskManager)_taskManager).jobs.Count == 3);
         }
         [TestMethod]
         public async Task TestGetTaskingNoTasks() {

diff --git a/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/TestClasses/TestTaskManager.cs b/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/TestClasses/TestTaskManager.cs
@@ -35,5 +35,10 @@ public async Task StartTaskAsync(ServerJob job)
             jobs.Add(job);
             return;
         }
+
+        public void WaitForNumberOfJobs(int numberOfJobs)
+        {
+            while(jobs.Count < numberOfJobs) { }
+        }
     }
 }