1  #include "pch.h"
  2  #include "CppUnitTest.h"
  3  #include "../CLILogic.h"
  4  #include <map>
  5  
  6  using namespace Microsoft::VisualStudio::CppUnitTestFramework;
  7  
  8  namespace FileLocksmithCLIUnitTests
  9  {
 10      struct MockProcessFinder : IProcessFinder
 11      {
 12          std::vector<ProcessResult> results;
 13          std::vector<ProcessResult> find(const std::vector<std::wstring>& paths) override
 14          {
 15              (void)paths;
 16              return results;
 17          }
 18      };
 19  
 20      struct MockProcessTerminator : IProcessTerminator
 21      {
 22          bool shouldSucceed = true;
 23          std::vector<DWORD> terminatedPids;
 24          bool terminate(DWORD pid) override
 25          {
 26              terminatedPids.push_back(pid);
 27              return shouldSucceed;
 28          }
 29      };
 30  
 31      struct MockStringProvider : IStringProvider
 32      {
 33          std::map<UINT, std::wstring> strings;
 34          std::wstring GetString(UINT id) override
 35          {
 36              if (strings.count(id)) return strings[id];
 37              return L"String_" + std::to_wstring(id);
 38          }
 39      };
 40  
 41      TEST_CLASS(CLITests)
 42      {
 43      public:
 44          
 45          TEST_METHOD(TestNoArgs)
 46          {
 47              MockProcessFinder finder;
 48              MockProcessTerminator terminator;
 49              MockStringProvider strings;
 50              
 51              wchar_t* argv[] = { (wchar_t*)L"exe" };
 52              auto result = run_command(1, argv, finder, terminator, strings);
 53              
 54              Assert::AreEqual(1, result.exit_code);
 55          }
 56  
 57          TEST_METHOD(TestHelp)
 58          {
 59              MockProcessFinder finder;
 60              MockProcessTerminator terminator;
 61              MockStringProvider strings;
 62              
 63              wchar_t* argv[] = { (wchar_t*)L"exe", (wchar_t*)L"--help" };
 64              auto result = run_command(2, argv, finder, terminator, strings);
 65              
 66              Assert::AreEqual(0, result.exit_code);
 67          }
 68  
 69          TEST_METHOD(TestFindProcesses)
 70          {
 71              MockProcessFinder finder;
 72              finder.results = { { L"process", 123, L"user", { L"file1" } } };
 73              MockProcessTerminator terminator;
 74              MockStringProvider strings;
 75              
 76              wchar_t* argv[] = { (wchar_t*)L"exe", (wchar_t*)L"file1" };
 77              auto result = run_command(2, argv, finder, terminator, strings);
 78              
 79              Assert::AreEqual(0, result.exit_code);
 80              Assert::IsTrue(result.output.find(L"123") != std::wstring::npos);
 81              Assert::IsTrue(result.output.find(L"process") != std::wstring::npos);
 82          }
 83  
 84          TEST_METHOD(TestJsonOutput)
 85          {
 86              MockProcessFinder finder;
 87              finder.results = { { L"process", 123, L"user", { L"file1" } } };
 88              MockProcessTerminator terminator;
 89              MockStringProvider strings;
 90              
 91              wchar_t* argv[] = { (wchar_t*)L"exe", (wchar_t*)L"file1", (wchar_t*)L"--json" };
 92              auto result = run_command(3, argv, finder, terminator, strings);
 93              
 94              Microsoft::VisualStudio::CppUnitTestFramework::Logger::WriteMessage(result.output.c_str());
 95  
 96              Assert::AreEqual(0, result.exit_code);
 97              Assert::IsTrue(result.output.find(L"\"pid\"") != std::wstring::npos);
 98              Assert::IsTrue(result.output.find(L"123") != std::wstring::npos);
 99          }
100  
101          TEST_METHOD(TestKill)
102          {
103              MockProcessFinder finder;
104              finder.results = { { L"process", 123, L"user", { L"file1" } } };
105              MockProcessTerminator terminator;
106              MockStringProvider strings;
107              
108              wchar_t* argv[] = { (wchar_t*)L"exe", (wchar_t*)L"file1", (wchar_t*)L"--kill" };
109              auto result = run_command(3, argv, finder, terminator, strings);
110              
111              Assert::AreEqual(0, result.exit_code);
112              Assert::AreEqual((size_t)1, terminator.terminatedPids.size());
113              Assert::AreEqual((DWORD)123, terminator.terminatedPids[0]);
114          }
115  
116          TEST_METHOD(TestTimeout)
117          {
118              MockProcessFinder finder;
119              // Always return results so it waits
120              finder.results = { { L"process", 123, L"user", { L"file1" } } };
121              MockProcessTerminator terminator;
122              MockStringProvider strings;
123              
124              wchar_t* argv[] = { (wchar_t*)L"exe", (wchar_t*)L"file1", (wchar_t*)L"--wait", (wchar_t*)L"--timeout", (wchar_t*)L"100" };
125              auto result = run_command(5, argv, finder, terminator, strings);
126              
127              Assert::AreEqual(1, result.exit_code);
128          }
129      };
130  }