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 index 01aaa866..ecc4800b 100644 --- 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 index 781a203c..daefdcd0 100644 --- 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) { } + } } }