Use open_memstream to override stdout during test execution.
The copy of the original stdout is held in env.stdout and used
to print subtest info and dump failed log.

test_{v,}printf are now simple wrappers around stdout and will be
removed in the next patch.

Cc: Andrii Nakryiko <andr...@fb.com>
Signed-off-by: Stanislav Fomichev <s...@google.com>
---
 tools/testing/selftests/bpf/test_progs.c | 100 ++++++++++-------------
 tools/testing/selftests/bpf/test_progs.h |   2 +-
 2 files changed, 46 insertions(+), 56 deletions(-)

diff --git a/tools/testing/selftests/bpf/test_progs.c 
b/tools/testing/selftests/bpf/test_progs.c
index db00196c8315..00d1565d01a3 100644
--- a/tools/testing/selftests/bpf/test_progs.c
+++ b/tools/testing/selftests/bpf/test_progs.c
@@ -40,14 +40,22 @@ static bool should_run(struct test_selector *sel, int num, 
const char *name)
 
 static void dump_test_log(const struct prog_test_def *test, bool failed)
 {
-       if (env.verbose || test->force_log || failed) {
-               if (env.log_cnt) {
-                       fprintf(stdout, "%s", env.log_buf);
-                       if (env.log_buf[env.log_cnt - 1] != '\n')
-                               fprintf(stdout, "\n");
+       if (stdout == env.stdout)
+               return;
+
+       fflush(stdout); /* exports env.log_buf & env.log_cap */
+
+       if (env.log_cap && (env.verbose || test->force_log || failed)) {
+               int len = strlen(env.log_buf);
+
+               if (len) {
+                       fprintf(env.stdout, "%s", env.log_buf);
+                       if (env.log_buf[len - 1] != '\n')
+                               fprintf(env.stdout, "\n");
+
+                       fseeko(stdout, 0, SEEK_SET); /* rewind */
                }
        }
-       env.log_cnt = 0;
 }
 
 void test__end_subtest()
@@ -62,7 +70,7 @@ void test__end_subtest()
 
        dump_test_log(test, sub_error_cnt);
 
-       printf("#%d/%d %s:%s\n",
+       fprintf(env.stdout, "#%d/%d %s:%s\n",
               test->test_num, test->subtest_num,
               test->subtest_name, sub_error_cnt ? "FAIL" : "OK");
 }
@@ -100,53 +108,7 @@ void test__force_log() {
 
 void test__vprintf(const char *fmt, va_list args)
 {
-       size_t rem_sz;
-       int ret = 0;
-
-       if (env.verbose || (env.test && env.test->force_log)) {
-               vfprintf(stderr, fmt, args);
-               return;
-       }
-
-try_again:
-       rem_sz = env.log_cap - env.log_cnt;
-       if (rem_sz) {
-               va_list ap;
-
-               va_copy(ap, args);
-               /* we reserved extra byte for \0 at the end */
-               ret = vsnprintf(env.log_buf + env.log_cnt, rem_sz + 1, fmt, ap);
-               va_end(ap);
-
-               if (ret < 0) {
-                       env.log_buf[env.log_cnt] = '\0';
-                       fprintf(stderr, "failed to log w/ fmt '%s'\n", fmt);
-                       return;
-               }
-       }
-
-       if (!rem_sz || ret > rem_sz) {
-               size_t new_sz = env.log_cap * 3 / 2;
-               char *new_buf;
-
-               if (new_sz < 4096)
-                       new_sz = 4096;
-               if (new_sz < ret + env.log_cnt)
-                       new_sz = ret + env.log_cnt;
-
-               /* +1 for guaranteed space for terminating \0 */
-               new_buf = realloc(env.log_buf, new_sz + 1);
-               if (!new_buf) {
-                       fprintf(stderr, "failed to realloc log buffer: %d\n",
-                               errno);
-                       return;
-               }
-               env.log_buf = new_buf;
-               env.log_cap = new_sz;
-               goto try_again;
-       }
-
-       env.log_cnt += ret;
+       vprintf(fmt, args);
 }
 
 void test__printf(const char *fmt, ...)
@@ -477,6 +439,32 @@ static error_t parse_arg(int key, char *arg, struct 
argp_state *state)
        return 0;
 }
 
+static void stdout_hijack(void)
+{
+       if (env.verbose || (env.test && env.test->force_log)) {
+               /* nothing to do, output to stdout by default */
+               return;
+       }
+
+       /* stdout -> buffer */
+       fflush(stdout);
+       stdout = open_memstream(&env.log_buf, &env.log_cap);
+}
+
+static void stdout_restore(void)
+{
+       if (stdout == env.stdout)
+               return;
+
+       fclose(stdout);
+       free(env.log_buf);
+
+       env.log_buf = NULL;
+       env.log_cap = 0;
+
+       stdout = env.stdout;
+}
+
 int main(int argc, char **argv)
 {
        static const struct argp argp = {
@@ -495,6 +483,7 @@ int main(int argc, char **argv)
        srand(time(NULL));
 
        env.jit_enabled = is_jit_enabled();
+       env.stdout = stdout;
 
        for (i = 0; i < prog_test_cnt; i++) {
                struct prog_test_def *test = &prog_test_defs[i];
@@ -508,6 +497,7 @@ int main(int argc, char **argv)
                                test->test_num, test->test_name))
                        continue;
 
+               stdout_hijack();
                test->run_test();
                /* ensure last sub-test is finalized properly */
                if (test->subtest_name)
@@ -522,6 +512,7 @@ int main(int argc, char **argv)
                        env.succ_cnt++;
 
                dump_test_log(test, test->error_cnt);
+               stdout_restore();
 
                printf("#%d %s:%s\n", test->test_num, test->test_name,
                       test->error_cnt ? "FAIL" : "OK");
@@ -529,7 +520,6 @@ int main(int argc, char **argv)
        printf("Summary: %d/%d PASSED, %d FAILED\n",
               env.succ_cnt, env.sub_succ_cnt, env.fail_cnt);
 
-       free(env.log_buf);
        free(env.test_selector.num_set);
        free(env.subtest_selector.num_set);
 
diff --git a/tools/testing/selftests/bpf/test_progs.h 
b/tools/testing/selftests/bpf/test_progs.h
index afd14962456f..9fd89078494f 100644
--- a/tools/testing/selftests/bpf/test_progs.h
+++ b/tools/testing/selftests/bpf/test_progs.h
@@ -56,8 +56,8 @@ struct test_env {
        bool jit_enabled;
 
        struct prog_test_def *test;
+       FILE *stdout;
        char *log_buf;
-       size_t log_cnt;
        size_t log_cap;
 
        int succ_cnt; /* successful tests */
-- 
2.22.0.770.g0f2c4a37fd-goog

Reply via email to