aboutsummaryrefslogtreecommitdiff
path: root/libexec/atf/atf-pytest-wrapper/atf_pytest_wrapper.cpp
blob: 6baa85999070966fc075febc11463f09e5f4d6c2 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
#include <format>
#include <iostream>
#include <map>
#include <string>
#include <vector>
#include <stdlib.h>
#include <unistd.h>

class Handler {
  private:
    const std::string kPytestName = "pytest";
    const std::string kCleanupSuffix = ":cleanup";
    const std::string kPythonPathEnv = "PYTHONPATH";
    const std::string kAtfVar = "_ATF_VAR_";
  public:
    // Test listing requested
    bool flag_list = false;
    // Output debug data (will break listing)
    bool flag_debug = false;
    // Cleanup for the test requested
    bool flag_cleanup = false;
    // Test source directory (provided by ATF)
    std::string src_dir;
    // Path to write test status to (provided by ATF)
    std::string dst_file;
    // Path to add to PYTHONPATH (provided by the schebang args)
    std::string python_path;
    // Path to the script (provided by the schebang wrapper)
    std::string script_path;
    // Name of the test to run (provided by ATF)
    std::string test_name;
    // kv pairs (provided by ATF)
    std::map<std::string,std::string> kv_map;
    // our binary name
    std::string binary_name;

    static std::vector<std::string> ToVector(int argc, char **argv) {
      std::vector<std::string> ret;

      for (int i = 0; i < argc; i++) {
        ret.emplace_back(std::string(argv[i]));
      }
      return ret;
    }

    static void PrintVector(std::string prefix, const std::vector<std::string> &vec) {
      std::cerr << prefix << ": ";
      for (auto &val: vec) {
        std::cerr << "'" << val << "' ";
      }
      std::cerr << std::endl;
    }

    void Usage(std::string msg, bool exit_with_error) {
      std::cerr << binary_name << ": ERROR: " << msg << "." << std::endl;
      std::cerr << binary_name << ": See atf-test-program(1) for usage details." << std::endl;
      exit(exit_with_error != 0);
    }

    // Parse args received from the OS. There can be multiple valid options:
    // * with schebang args (#!/binary -P/path):
    // atf_wrap '-P /path' /path/to/script -l
    // * without schebang args
    // atf_wrap /path/to/script -l
    // Running test:
    // atf_wrap '-P /path' /path/to/script -r /path1 -s /path2 -vk1=v1 testname 
    void Parse(int argc, char **argv) {
      if (flag_debug) {
        PrintVector("IN", ToVector(argc, argv));
      }
      // getopt() skips the first argument (as it is typically binary name)
      // it is possible to have either '-P\s*/path' followed by the script name
      // or just the script name. Parse kernel-provided arg manually and adjust
      // array to make getopt work

      binary_name = std::string(argv[0]);
      argc--; argv++;
      // parse -P\s*path from the kernel.
      if (argc > 0 && !strncmp(argv[0], "-P", 2)) {
        char *path = &argv[0][2];
        while (*path == ' ')
          path++;
        python_path = std::string(path);
        argc--; argv++;
      }

      // The next argument is a script name. Copy and keep argc/argv the same
      // Show usage for empty args
      if (argc == 0) {
          Usage("Must provide a test case name", true);
      }
      script_path = std::string(argv[0]);

      int c;
      while ((c = getopt(argc, argv, "lr:s:v:")) != -1) {
        switch (c) {
        case 'l':
          flag_list = true;
          break;
        case 's':
          src_dir = std::string(optarg);
          break;
        case 'r':
          dst_file = std::string(optarg);
          break;
        case 'v':
	  {
	    std::string kv = std::string(optarg);
	    size_t splitter = kv.find("=");
	    if (splitter == std::string::npos) {
	      Usage("Unknown variable: " + kv, true);
	    }
	    kv_map[kv.substr(0, splitter)] = kv.substr(splitter + 1);
	  }
          break;
        default:
          Usage("Unknown option -" + std::string(1, static_cast<char>(c)), true);
        }
      }
      argc -= optind;
      argv += optind;

      if (flag_list) {
        return;
      }
      // There should be just one argument with the test name
      if (argc != 1) {
        Usage("Must provide a test case name", true);
      }
      test_name = std::string(argv[0]);
      if (test_name.size() > kCleanupSuffix.size() &&
          std::equal(kCleanupSuffix.rbegin(), kCleanupSuffix.rend(), test_name.rbegin())) {
        test_name = test_name.substr(0, test_name.size() - kCleanupSuffix.size());
        flag_cleanup = true;
      }
    }

    std::vector<std::string> BuildArgs() {
      std::vector<std::string> args = {"pytest", "-p", "no:cacheprovider", "-s", "--atf"};

      if (flag_list) {
        args.push_back("--co");
        args.push_back(script_path);
        return args;
      }
      if (flag_cleanup) {
        args.push_back("--atf-cleanup");
      }
      // workaround pytest parser bug:
      // https://github.com/pytest-dev/pytest/issues/3097
      // use '--arg=value' format instead of '--arg value' for all
      // path-like options
      if (!src_dir.empty()) {
        args.push_back("--atf-source-dir=" + src_dir);
      }
      if (!dst_file.empty()) {
        args.push_back("--atf-file=" + dst_file);
      }
      // Create nodeid from the test path &name
      args.push_back(script_path + "::" + test_name);
      return args;
    }

    void SetPythonPath() {
      if (!python_path.empty()) {
        char *env_path = getenv(kPythonPathEnv.c_str());
        if (env_path != nullptr) {
          python_path = python_path + ":" + std::string(env_path);
        }
        setenv(kPythonPathEnv.c_str(), python_path.c_str(), 1);
      }
    }

    void SetEnv() {
      SetPythonPath();

      // Pass ATF kv pairs as env variables to avoid dealing with
      // pytest parser
      for (auto [k, v]: kv_map) {
	setenv((kAtfVar + k).c_str(), v.c_str(), 1);
      }
    }

    int Run(std::string binary, std::vector<std::string> args) {
      if (flag_debug) {
        PrintVector("OUT", args);
      }
      // allocate array with final NULL
      char **arr = new char*[args.size() + 1]();
      for (unsigned long i = 0; i < args.size(); i++) {
	// work around 'char *const *'
        arr[i] = strdup(args[i].c_str());
      }
      return (execvp(binary.c_str(), arr) != 0);
    }

    int Process() {
      SetEnv();
      return Run(kPytestName, BuildArgs());
    }
};


int main(int argc, char **argv) {
  Handler handler;

  handler.Parse(argc, argv);
  return handler.Process();
}