From 3488339d5119b46303d357202aef0db007d466b6 Mon Sep 17 00:00:00 2001
From: Deukhoofd <Deukhoofd@gmail.com>
Date: Sun, 22 Aug 2021 19:29:00 +0200
Subject: [PATCH] Fixes several memory issues.

---
 src/Tester/Test.hpp          |  9 +++++----
 src/Tester/TestFunctions.hpp | 22 +++++++++++++++++++---
 src/Tester/TestRunner.hpp    | 32 +++++++++++++++++++++++++-------
 src/main.cpp                 |  6 ++----
 4 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/src/Tester/Test.hpp b/src/Tester/Test.hpp
index cf35c36..3edd8ff 100644
--- a/src/Tester/Test.hpp
+++ b/src/Tester/Test.hpp
@@ -11,16 +11,17 @@ public:
         : _name(name), _function(function), _env(new TestEnvironment()) {}
 
     void Run(asIScriptContext* ctx) {
+        ctx->PushState();
         ctx->Prepare(_function);
-        ctx->SetUserData(_env);
+        ctx->SetUserData(_env.get());
         auto e = ctx->Execute();
         if (e == asEXECUTION_EXCEPTION) {
             _errorMessage = ctx->GetExceptionString();
             _result = TestResult::Failed;
-            ctx->Release();
+            ctx->PopState();
             return;
         }
-        ctx->Release();
+        ctx->PopState();
         Ensure(e == asEXECUTION_FINISHED);
         _result = TestResult::Success;
     }
@@ -31,7 +32,7 @@ public:
 private:
     std::string _name;
     asIScriptFunction* _function;
-    TestEnvironment* _env;
+    std::unique_ptr<TestEnvironment> _env;
 
     TestResult _result = TestResult::NotRan;
     std::string _errorMessage = "";
diff --git a/src/Tester/TestFunctions.hpp b/src/Tester/TestFunctions.hpp
index e0ebff4..ed827d2 100644
--- a/src/Tester/TestFunctions.hpp
+++ b/src/Tester/TestFunctions.hpp
@@ -40,7 +40,7 @@ private:
         eqMsg = ss.str();                                                                                              \
     }
 
-    static bool AssertEquals(i32 expected, i32 actual) {
+    static bool AssertEqualsI32(i32 expected, i32 actual) {
         if (expected != actual) {
             auto data = GetAssertionData();
             GetEqualityMessage(expected, actual);
@@ -49,12 +49,28 @@ private:
         return true;
     }
 
+    static bool AssertEqualsString(const std::string& expected, const std::string& actual) {
+        if (expected != actual) {
+            auto data = GetAssertionData();
+            std::string eqMsg;
+            {
+                std::ostringstream ss;
+                ss << "Expected: '" << expected << "', but was: '" << actual << "'";
+                eqMsg = ss.str();
+            }
+            throw AssertionFailed(data.FileName, data.Line, eqMsg);
+        }
+        return true;
+    }
+
 public:
     static void Register(AngelScriptResolver* scriptResolver) {
         auto engine = scriptResolver->GetBuilder().GetEngine();
         Ensure(engine->RegisterGlobalFunction("bool Assert(bool expression)", asFUNCTION(Assert), asCALL_CDECL) >= 0);
-        Ensure(engine->RegisterGlobalFunction("bool AssertEquals(int expected, int actual)", asFUNCTION(AssertEquals),
-                                              asCALL_CDECL) >= 0);
+        Ensure(engine->RegisterGlobalFunction("bool AssertEquals(int expected, int actual)",
+                                              asFUNCTION(AssertEqualsI32), asCALL_CDECL) >= 0);
+        Ensure(engine->RegisterGlobalFunction("bool AssertEquals(const string &in expected, const string &in actual)",
+                                              asFUNCTION(AssertEqualsString), asCALL_CDECL) >= 0);
     }
 };
 
diff --git a/src/Tester/TestRunner.hpp b/src/Tester/TestRunner.hpp
index a96d604..b3f8ee7 100644
--- a/src/Tester/TestRunner.hpp
+++ b/src/Tester/TestRunner.hpp
@@ -5,7 +5,7 @@
 #include "Test.hpp"
 
 class TestRunner {
-    ArbUt::Dictionary<std::string, Test> _tests;
+    ArbUt::Dictionary<std::string, std::unique_ptr<Test>> _tests;
 
 public:
     TestRunner(AngelScriptResolver* scriptResolver) {
@@ -18,24 +18,42 @@ public:
                 auto meta = AngelscriptMetadata(m);
                 if (meta.GetIdentifier() == "Test"_cnc) {
                     auto name = meta.GetParameter("name"_cnc);
-                    _tests.Insert(name, Test(name, func));
+                    _tests.GetStdMap().insert({name, std::make_unique<Test>(name, func)});
                 }
             }
         }
     }
 
-    void RunAll(asIScriptEngine* engine) {
+    i32 RunAll(asIScriptEngine* engine) {
         auto ctx = engine->CreateContext();
         for (auto& test : _tests) {
-            test.second.Run(ctx);
+            test.second->Run(ctx);
         }
+        ctx->Release();
 
+        size_t notRunTests = 0;
+        size_t successfulTests = 0;
+        size_t failedTests = 0;
         for (auto& test : _tests) {
-            if (test.second.GetResult() == TestResult::Failed) {
-                std::cout << "Test '" << test.first << "' failed with message: " << std::endl;
-                std::cout << test.second.GetErrorMessage() << std::endl << std::endl;
+            auto result = test.second->GetResult();
+            switch (result) {
+                case TestResult::NotRan: notRunTests += 1; break;
+                case TestResult::Success: successfulTests += 1; break;
+                case TestResult::Failed: {
+                    std::cout << "Test '" << test.first << "' failed with message: " << std::endl;
+                    std::cout << test.second->GetErrorMessage() << std::endl << std::endl;
+                    failedTests += 1;
+                    break;
+                }
             }
         }
+
+        std::cout << "Ran " << successfulTests + failedTests << " tests. Passed tests: " << successfulTests
+                  << ". Failed tests: " << failedTests << "." << std::endl;
+        if (failedTests > 0) {
+            return 1;
+        }
+        return 0;
     }
 };
 
diff --git a/src/main.cpp b/src/main.cpp
index 5579fd9..db3d605 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -49,8 +49,6 @@ int main(int argc, char** argv) {
     }
     auto* scriptResolver = dynamic_cast<AngelScriptResolver*>(battleLib.GetValue()->GetScriptResolver().get());
     auto testRunner = TestRunner(scriptResolver);
-    auto engine = scriptResolver->GetBuilder().GetEngine();
-    testRunner.RunAll(engine);
-
-    return 0;
+    auto* engine = scriptResolver->GetBuilder().GetEngine();
+    return testRunner.RunAll(engine);
 }
\ No newline at end of file