1  #pragma once
  2  
  3  #include "pch.h"
  4  #include <string>
  5  #include <filesystem>
  6  #include <fstream>
  7  #include <random>
  8  
  9  namespace TestHelpers
 10  {
 11      // RAII helper for creating and cleaning up temporary files
 12      class TempFile
 13      {
 14      public:
 15          TempFile(const std::wstring& content = L"", const std::wstring& extension = L".txt")
 16          {
 17              wchar_t tempPath[MAX_PATH];
 18              GetTempPathW(MAX_PATH, tempPath);
 19  
 20              // Generate a unique filename
 21              std::random_device rd;
 22              std::mt19937 gen(rd());
 23              std::uniform_int_distribution<> dis(10000, 99999);
 24  
 25              m_path = std::wstring(tempPath) + L"test_" + std::to_wstring(dis(gen)) + extension;
 26  
 27              if (!content.empty())
 28              {
 29                  std::wofstream file(m_path);
 30                  file << content;
 31              }
 32          }
 33  
 34          ~TempFile()
 35          {
 36              if (std::filesystem::exists(m_path))
 37              {
 38                  std::filesystem::remove(m_path);
 39              }
 40          }
 41  
 42          TempFile(const TempFile&) = delete;
 43          TempFile& operator=(const TempFile&) = delete;
 44  
 45          const std::wstring& path() const { return m_path; }
 46  
 47          void write(const std::string& content)
 48          {
 49              std::ofstream file(m_path, std::ios::binary);
 50              file << content;
 51          }
 52  
 53          void write(const std::wstring& content)
 54          {
 55              std::wofstream file(m_path);
 56              file << content;
 57          }
 58  
 59          std::wstring read()
 60          {
 61              std::wifstream file(m_path);
 62              return std::wstring((std::istreambuf_iterator<wchar_t>(file)),
 63                                 std::istreambuf_iterator<wchar_t>());
 64          }
 65  
 66      private:
 67          std::wstring m_path;
 68      };
 69  
 70      // RAII helper for creating and cleaning up temporary directories
 71      class TempDirectory
 72      {
 73      public:
 74          TempDirectory()
 75          {
 76              wchar_t tempPath[MAX_PATH];
 77              GetTempPathW(MAX_PATH, tempPath);
 78  
 79              std::random_device rd;
 80              std::mt19937 gen(rd());
 81              std::uniform_int_distribution<> dis(10000, 99999);
 82  
 83              m_path = std::wstring(tempPath) + L"testdir_" + std::to_wstring(dis(gen));
 84              std::filesystem::create_directories(m_path);
 85          }
 86  
 87          ~TempDirectory()
 88          {
 89              if (std::filesystem::exists(m_path))
 90              {
 91                  std::filesystem::remove_all(m_path);
 92              }
 93          }
 94  
 95          TempDirectory(const TempDirectory&) = delete;
 96          TempDirectory& operator=(const TempDirectory&) = delete;
 97  
 98          const std::wstring& path() const { return m_path; }
 99  
100      private:
101          std::wstring m_path;
102      };
103  
104      // Registry test key path - use HKCU for non-elevated tests
105      inline const std::wstring TestRegistryPath = L"Software\\PowerToys\\UnitTests";
106  
107      // RAII helper for registry key creation/cleanup
108      class TestRegistryKey
109      {
110      public:
111          TestRegistryKey(const std::wstring& subKey = L"")
112          {
113              m_path = TestRegistryPath;
114              if (!subKey.empty())
115              {
116                  m_path += L"\\" + subKey;
117              }
118  
119              HKEY key;
120              if (RegCreateKeyExW(HKEY_CURRENT_USER, m_path.c_str(), 0, nullptr,
121                                 REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, nullptr,
122                                 &key, nullptr) == ERROR_SUCCESS)
123              {
124                  RegCloseKey(key);
125                  m_created = true;
126              }
127          }
128  
129          ~TestRegistryKey()
130          {
131              if (m_created)
132              {
133                  RegDeleteTreeW(HKEY_CURRENT_USER, m_path.c_str());
134              }
135          }
136  
137          TestRegistryKey(const TestRegistryKey&) = delete;
138          TestRegistryKey& operator=(const TestRegistryKey&) = delete;
139  
140          bool isValid() const { return m_created; }
141          const std::wstring& path() const { return m_path; }
142  
143          bool setStringValue(const std::wstring& name, const std::wstring& value)
144          {
145              HKEY key;
146              if (RegOpenKeyExW(HKEY_CURRENT_USER, m_path.c_str(), 0, KEY_SET_VALUE, &key) != ERROR_SUCCESS)
147              {
148                  return false;
149              }
150  
151              auto result = RegSetValueExW(key, name.c_str(), 0, REG_SZ,
152                                          reinterpret_cast<const BYTE*>(value.c_str()),
153                                          static_cast<DWORD>((value.length() + 1) * sizeof(wchar_t)));
154              RegCloseKey(key);
155              return result == ERROR_SUCCESS;
156          }
157  
158          bool setDwordValue(const std::wstring& name, DWORD value)
159          {
160              HKEY key;
161              if (RegOpenKeyExW(HKEY_CURRENT_USER, m_path.c_str(), 0, KEY_SET_VALUE, &key) != ERROR_SUCCESS)
162              {
163                  return false;
164              }
165  
166              auto result = RegSetValueExW(key, name.c_str(), 0, REG_DWORD,
167                                          reinterpret_cast<const BYTE*>(&value), sizeof(DWORD));
168              RegCloseKey(key);
169              return result == ERROR_SUCCESS;
170          }
171  
172      private:
173          std::wstring m_path;
174          bool m_created = false;
175      };
176  
177      // Helper to wait for a condition with timeout
178      template<typename Predicate>
179      bool WaitFor(Predicate pred, std::chrono::milliseconds timeout = std::chrono::milliseconds(5000))
180      {
181          auto start = std::chrono::steady_clock::now();
182          while (!pred())
183          {
184              if (std::chrono::steady_clock::now() - start > timeout)
185              {
186                  return false;
187              }
188              std::this_thread::sleep_for(std::chrono::milliseconds(10));
189          }
190          return true;
191      }
192  }