1 // vim: ts=2 sw=2 et
2
3 #include <format>
4 #include <iostream>
5 #include <map>
6 #include <string>
7 #include <vector>
8 #include <stdlib.h>
9 #include <unistd.h>
10
11 class Handler {
12 private:
13 const std::string kPytestName = "pytest";
14 const std::string kCleanupSuffix = ":cleanup";
15 const std::string kPythonPathEnv = "PYTHONPATH";
16 const std::string kAtfVar = "_ATF_VAR_";
17 public:
18 // Test listing requested
19 bool flag_list = false;
20 // Output debug data (will break listing)
21 bool flag_debug = false;
22 // Cleanup for the test requested
23 bool flag_cleanup = false;
24 // Test source directory (provided by ATF)
25 std::string src_dir;
26 // Path to write test status to (provided by ATF)
27 std::string dst_file;
28 // Path to add to PYTHONPATH (provided by the schebang args)
29 std::string python_path;
30 // Path to the script (provided by the schebang wrapper)
31 std::string script_path;
32 // Name of the test to run (provided by ATF)
33 std::string test_name;
34 // kv pairs (provided by ATF)
35 std::map<std::string,std::string> kv_map;
36 // our binary name
37 std::string binary_name;
38
ToVector(int argc,char ** argv)39 static std::vector<std::string> ToVector(int argc, char **argv) {
40 std::vector<std::string> ret;
41
42 for (int i = 0; i < argc; i++) {
43 ret.emplace_back(std::string(argv[i]));
44 }
45 return ret;
46 }
47
PrintVector(std::string prefix,const std::vector<std::string> & vec)48 static void PrintVector(std::string prefix, const std::vector<std::string> &vec) {
49 std::cerr << prefix << ": ";
50 for (auto &val: vec) {
51 std::cerr << "'" << val << "' ";
52 }
53 std::cerr << std::endl;
54 }
55
Usage(std::string msg,bool exit_with_error)56 void Usage(std::string msg, bool exit_with_error) {
57 std::cerr << binary_name << ": ERROR: " << msg << "." << std::endl;
58 std::cerr << binary_name << ": See atf-test-program(1) for usage details." << std::endl;
59 exit(exit_with_error != 0);
60 }
61
62 // Parse args received from the OS. There can be multiple valid options:
63 // * with schebang args (#!/binary -P/path):
64 // atf_wrap '-P /path' /path/to/script -l
65 // * without schebang args
66 // atf_wrap /path/to/script -l
67 // Running test:
68 // atf_wrap '-P /path' /path/to/script -r /path1 -s /path2 -vk1=v1 testname
Parse(int argc,char ** argv)69 void Parse(int argc, char **argv) {
70 if (flag_debug) {
71 PrintVector("IN", ToVector(argc, argv));
72 }
73 // getopt() skips the first argument (as it is typically binary name)
74 // it is possible to have either '-P\s*/path' followed by the script name
75 // or just the script name. Parse kernel-provided arg manually and adjust
76 // array to make getopt work
77
78 binary_name = std::string(argv[0]);
79 argc--; argv++;
80 // parse -P\s*path from the kernel.
81 if (argc > 0 && !strncmp(argv[0], "-P", 2)) {
82 char *path = &argv[0][2];
83 while (*path == ' ')
84 path++;
85 python_path = std::string(path);
86 argc--; argv++;
87 }
88
89 // The next argument is a script name. Copy and keep argc/argv the same
90 // Show usage for empty args
91 if (argc == 0) {
92 Usage("Must provide a test case name", true);
93 }
94 script_path = std::string(argv[0]);
95
96 int c;
97 while ((c = getopt(argc, argv, "lr:s:v:")) != -1) {
98 switch (c) {
99 case 'l':
100 flag_list = true;
101 break;
102 case 's':
103 src_dir = std::string(optarg);
104 break;
105 case 'r':
106 dst_file = std::string(optarg);
107 break;
108 case 'v':
109 {
110 std::string kv = std::string(optarg);
111 size_t splitter = kv.find("=");
112 if (splitter == std::string::npos) {
113 Usage("Unknown variable: " + kv, true);
114 }
115 kv_map[kv.substr(0, splitter)] = kv.substr(splitter + 1);
116 }
117 break;
118 default:
119 Usage("Unknown option -" + std::string(1, static_cast<char>(c)), true);
120 }
121 }
122 argc -= optind;
123 argv += optind;
124
125 if (flag_list) {
126 return;
127 }
128 // There should be just one argument with the test name
129 if (argc != 1) {
130 Usage("Must provide a test case name", true);
131 }
132 test_name = std::string(argv[0]);
133 if (test_name.size() > kCleanupSuffix.size() &&
134 std::equal(kCleanupSuffix.rbegin(), kCleanupSuffix.rend(), test_name.rbegin())) {
135 test_name = test_name.substr(0, test_name.size() - kCleanupSuffix.size());
136 flag_cleanup = true;
137 }
138 }
139
BuildArgs()140 std::vector<std::string> BuildArgs() {
141 std::vector<std::string> args = {"pytest", "-vv", "-p",
142 "no:cacheprovider", "-s", "--atf"};
143
144 args.push_back("--confcutdir=" + python_path);
145
146 if (flag_list) {
147 args.push_back("--co");
148 args.push_back(script_path);
149 return args;
150 }
151 if (flag_cleanup) {
152 args.push_back("--atf-cleanup");
153 }
154 // workaround pytest parser bug:
155 // https://github.com/pytest-dev/pytest/issues/3097
156 // use '--arg=value' format instead of '--arg value' for all
157 // path-like options
158 if (!src_dir.empty()) {
159 args.push_back("--atf-source-dir=" + src_dir);
160 }
161 if (!dst_file.empty()) {
162 args.push_back("--atf-file=" + dst_file);
163 }
164 // Create nodeid from the test path &name
165 args.push_back(script_path + "::" + test_name);
166 return args;
167 }
168
SetPythonPath()169 void SetPythonPath() {
170 if (!python_path.empty()) {
171 char *env_path = getenv(kPythonPathEnv.c_str());
172 if (env_path != nullptr) {
173 python_path = python_path + ":" + std::string(env_path);
174 }
175 setenv(kPythonPathEnv.c_str(), python_path.c_str(), 1);
176 }
177 }
178
SetEnv()179 void SetEnv() {
180 SetPythonPath();
181
182 // Pass ATF kv pairs as env variables to avoid dealing with
183 // pytest parser
184 for (auto [k, v]: kv_map) {
185 setenv((kAtfVar + k).c_str(), v.c_str(), 1);
186 }
187 }
188
Run(std::string binary,std::vector<std::string> args)189 bool Run(std::string binary, std::vector<std::string> args) {
190 if (flag_debug) {
191 PrintVector("OUT", args);
192 }
193 // allocate array with final NULL
194 char **arr = new char*[args.size() + 1]();
195 for (unsigned long i = 0; i < args.size(); i++) {
196 // work around 'char *const *'
197 arr[i] = strdup(args[i].c_str());
198 }
199 return execvp(binary.c_str(), arr) == 0;
200 }
201
ReportError()202 void ReportError() {
203 if (flag_list) {
204 std::cout << "Content-Type: application/X-atf-tp; version=\"1\"";
205 std::cout << std::endl << std::endl;
206 std::cout << "ident: __test_cases_list_"<< kPytestName << "_binary_" <<
207 "not_found__" << std::endl;
208 } else {
209 std::cout << "execvp(" << kPytestName << ") failed: " <<
210 std::strerror(errno) << std::endl;
211 }
212 }
213
Process()214 int Process() {
215 SetEnv();
216 if (!Run(kPytestName, BuildArgs())) {
217 ReportError();
218 }
219 return 0;
220 }
221 };
222
223
main(int argc,char ** argv)224 int main(int argc, char **argv) {
225 Handler handler;
226
227 handler.Parse(argc, argv);
228 return handler.Process();
229 }
230