Line data Source code
1 : // SPDX-FileCopyrightText: 2025 PairInteraction Developers
2 : // SPDX-License-Identifier: LGPL-3.0-or-later
3 :
4 : #include "pairinteraction/database/ParquetManager.hpp"
5 :
6 : #include "pairinteraction/database/Database.hpp"
7 : #include "pairinteraction/database/GitHubDownloader.hpp"
8 :
9 : #include <doctest/doctest.h>
10 : #include <duckdb.hpp>
11 : #include <filesystem>
12 : #include <fmt/format.h>
13 : #include <fstream>
14 : #include <memory>
15 : #include <miniz.h>
16 : #include <nlohmann/json.hpp>
17 :
18 : namespace pairinteraction {
19 : class MockDownloader : public GitHubDownloader {
20 : public:
21 : std::future<GitHubDownloader::Result>
22 12 : download(const std::string &remote_url, const std::string & /*if_none_match*/ = "",
23 : bool /*use_octet_stream*/ = false) const override {
24 12 : GitHubDownloader::Result result;
25 12 : result.status_code = 200;
26 12 : result.rate_limit.remaining = 60;
27 12 : result.rate_limit.reset_time = 2147483647;
28 :
29 12 : if (remote_url == "/test/repo/path") {
30 : // This is a repo path request for a single release, return JSON with assets
31 2 : result.body = make_release("2.1").dump();
32 11 : } else if (remote_url == "/test/repo/releases") {
33 : // This is a repo path request for a list of releases. The latest release only
34 : // provides tables whose major version differs from COMPATIBLE_DATABASE_VERSION_MAJOR
35 : // and is thus incompatible, the older ones are compatible. In addition, the oldest
36 : // release contains an asset that is not provided anymore by newer releases.
37 2 : nlohmann::json releases = nlohmann::json::array();
38 4 : releases.push_back(make_release("3.0"));
39 4 : releases.push_back(make_release("2.1"));
40 6 : releases.push_back(make_release("2.0", {"misc", "retired"}));
41 2 : result.body = releases.dump();
42 11 : } else if (remote_url == "/test/repo/releases_unpublished") {
43 : // This is a repo path request for a list of releases whose latest entries are a draft
44 : // and a pre-release. They must be skipped although they provide compatible tables. The
45 : // draft does not contain any assets at all, as it is the case for drafts on GitHub
46 : // that have not been populated yet.
47 1 : nlohmann::json draft;
48 1 : draft["draft"] = true;
49 :
50 2 : nlohmann::json prerelease = make_release("2.4");
51 1 : prerelease["prerelease"] = true;
52 :
53 1 : nlohmann::json releases = nlohmann::json::array();
54 1 : releases.push_back(draft);
55 1 : releases.push_back(prerelease);
56 2 : releases.push_back(make_release("2.1"));
57 1 : result.body = releases.dump();
58 9 : } else if (remote_url == "/test/repo/releases_invalid") {
59 : // This is a repo path request whose response is not valid JSON
60 1 : result.body = "not a json";
61 7 : } else if (remote_url == "/rate_limit") {
62 : // This is the rate limit request
63 4 : result.body = "";
64 : } else {
65 : // This is the file download request
66 3 : std::string content = "updated_file_content";
67 3 : std::string filename = "misc_v2.1/wigner.parquet";
68 :
69 3 : mz_zip_archive zip_archive{};
70 3 : size_t zip_size = 0;
71 3 : void *zip_data = nullptr;
72 :
73 3 : mz_zip_writer_init_heap(&zip_archive, 0, 0);
74 3 : mz_zip_writer_add_mem(&zip_archive, filename.c_str(), content.data(), content.size(),
75 : MZ_BEST_SPEED);
76 3 : mz_zip_writer_finalize_heap_archive(&zip_archive, &zip_data, &zip_size);
77 :
78 3 : result.body = std::string(static_cast<char *>(zip_data), zip_size);
79 :
80 3 : mz_free(zip_data);
81 3 : mz_zip_writer_end(&zip_archive);
82 3 : }
83 :
84 48 : return std::async(std::launch::deferred, [result]() { return result; });
85 37 : }
86 :
87 : private:
88 : // Construct a release that provides the tables of the given version for the given assets
89 9 : static nlohmann::json make_release(const std::string &version,
90 : const std::vector<std::string> &keys = {"misc"}) {
91 9 : nlohmann::json assets = nlohmann::json::array();
92 20 : for (const auto &key : keys) {
93 11 : nlohmann::json asset;
94 22 : asset["name"] = fmt::format("{}_v{}.zip", key, version);
95 22 : asset["url"] = fmt::format("https://api.github.com/test/path/{}_v{}.zip", key, version);
96 11 : assets.push_back(asset);
97 11 : }
98 :
99 9 : nlohmann::json release;
100 9 : release["assets"] = assets;
101 18 : return release;
102 9 : }
103 : };
104 :
105 6 : TEST_CASE("ParquetManager functionality with mocked downloader") {
106 6 : MockDownloader downloader;
107 6 : auto test_dir = std::filesystem::temp_directory_path() / "pairinteraction_test_db";
108 6 : duckdb::DuckDB db(nullptr);
109 6 : duckdb::Connection con(db);
110 :
111 : // Create a manager for the given repository paths and scan the local and remote tables
112 6 : auto make_manager = [&](std::vector<std::string> repo_paths) {
113 12 : auto manager = std::make_unique<ParquetManager>(test_dir, downloader, std::move(repo_paths),
114 6 : con, false);
115 6 : manager->scan_local();
116 6 : manager->scan_remote();
117 6 : return manager;
118 0 : };
119 :
120 : // Construct the path to the wigner table of the misc asset of the given version
121 19 : auto wigner_path = [&](const std::string &version) {
122 57 : return (test_dir / "tables" / fmt::format("misc_v{}", version) / "wigner.parquet").string();
123 6 : };
124 :
125 : // Create a local wigner table of the misc asset of the given version
126 13 : auto add_local_table = [&](const std::string &version) {
127 26 : std::filesystem::create_directories(test_dir / "tables" / fmt::format("misc_v{}", version));
128 13 : std::ofstream(wigner_path(version)).close();
129 13 : };
130 :
131 : // Provide a compatible version and an incompatible version with a higher minor version
132 6 : add_local_table("2.0");
133 6 : add_local_table("1.4");
134 :
135 6 : SUBCASE("Check missing table") {
136 1 : auto manager = make_manager({});
137 :
138 5 : CHECK_THROWS_WITH_AS(manager->get_path("misc", "missing_table"),
139 : "No table 'missing_table.parquet' found for species 'misc'. The "
140 : "tables for the species are incomplete.",
141 : std::runtime_error);
142 7 : }
143 :
144 6 : SUBCASE("Check version parsing") {
145 : // Provide a second compatible version in addition
146 1 : add_local_table("2.2");
147 :
148 1 : auto manager = make_manager({});
149 :
150 : // The newest compatible version must be used, neither the oldest compatible version nor
151 : // the incompatible version with the highest minor version
152 1 : CHECK(manager->get_path("misc", "wigner") == wigner_path("2.2"));
153 7 : }
154 :
155 6 : SUBCASE("Check update table") {
156 2 : auto manager = make_manager({"/test/repo/path"});
157 :
158 1 : CHECK(manager->get_path("misc", "wigner") == wigner_path("2.1"));
159 :
160 1 : std::ifstream in(wigner_path("2.1"), std::ios::binary);
161 1 : std::stringstream buffer;
162 1 : buffer << in.rdbuf();
163 1 : CHECK(buffer.str() == "updated_file_content");
164 7 : }
165 :
166 6 : SUBCASE("Check update table if the latest release is incompatible") {
167 2 : auto manager = make_manager({"/test/repo/releases"});
168 :
169 : // The latest compatible release must be used, not the latest release
170 1 : CHECK(manager->get_path("misc", "wigner") == wigner_path("2.1"));
171 :
172 : // Assets that are only provided by older releases must not be used
173 5 : CHECK_THROWS_WITH_AS(
174 : manager->get_path("retired", "wigner"),
175 : "No tables found for species 'retired'. Check the spelling of the species.",
176 : std::runtime_error);
177 7 : }
178 :
179 6 : SUBCASE("Check invalid response of a repository") {
180 : // If one repository returns an invalid response, the download of database tables must be
181 : // disabled altogether so that the local tables are used instead of silently providing an
182 : // incomplete set of tables
183 3 : auto manager = make_manager({"/test/repo/releases", "/test/repo/releases_invalid"});
184 :
185 1 : CHECK(manager->get_path("misc", "wigner") == wigner_path("2.0"));
186 7 : }
187 :
188 6 : SUBCASE("Check update table if the latest releases are unpublished") {
189 2 : auto manager = make_manager({"/test/repo/releases_unpublished"});
190 :
191 : // The latest published release must be used, neither the draft nor the pre-release
192 1 : CHECK(manager->get_path("misc", "wigner") == wigner_path("2.1"));
193 7 : }
194 :
195 6 : std::filesystem::remove_all(test_dir);
196 14 : }
197 :
198 1 : DOCTEST_TEST_CASE("ParquetManager functionality with GitHub downloader") {
199 1 : if (!Database::get_global_instance().get_download_missing()) {
200 1 : DOCTEST_MESSAGE("Skipping test because download_missing is false.");
201 1 : return;
202 : }
203 0 : GitHubDownloader downloader;
204 0 : duckdb::DuckDB db(nullptr);
205 0 : duckdb::Connection con(db);
206 :
207 : std::vector<std::string> repo_paths = {
208 : "/repos/pairinteraction/database-sqdt/releases?per_page=100",
209 0 : "/repos/pairinteraction/database-mqdt/releases?per_page=100"};
210 0 : ParquetManager manager(Database::get_global_instance().get_database_dir(), downloader,
211 0 : repo_paths, con, Database::get_global_instance().get_use_cache());
212 0 : manager.scan_local();
213 0 : manager.scan_remote();
214 :
215 0 : std::string info = manager.get_versions_info();
216 :
217 : // Check that all species are present
218 : std::vector<std::string> should_contain = {
219 : "Cs", "K", "Li", "Na", "Rb", "Sr87_mqdt",
220 0 : "Sr88_sqdt", "Sr88_mqdt", "Yb171_mqdt", "Yb173_mqdt", "Yb174_mqdt", "misc"};
221 0 : for (const auto &substr : should_contain) {
222 0 : DOCTEST_CHECK(info.find(substr) != std::string::npos);
223 : }
224 0 : }
225 :
226 : } // namespace pairinteraction
|
