aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/platform/test_main.cc
blob: 11230c3f7bfc1ad86ec271cf9f7027d03de2a931 (plain)
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
// A program with a main that is suitable for unittests, including those
// that also define microbenchmarks.  Based on whether the user specified
// the --benchmark_filter flag which specifies which benchmarks to run,
// we will either run benchmarks or run the gtest tests in the program.

#include <iostream>

#include "tensorflow/core/platform/port.h"

#if defined(PLATFORM_GOOGLE) || defined(PLATFORM_POSIX_ANDROID) || \
    defined(PLATFORM_GOOGLE_ANDROID)
// main() is supplied by gunit_main
#else
#include "gtest/gtest.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/platform/test_benchmark.h"

GTEST_API_ int main(int argc, char** argv) {
  std::cout << "Running main() from test_main.cc\n";

  testing::InitGoogleTest(&argc, argv);
  for (int i = 1; i < argc; i++) {
    if (tensorflow::StringPiece(argv[i]).starts_with("--benchmarks=")) {
      const char* pattern = argv[i] + strlen("--benchmarks=");
      tensorflow::testing::Benchmark::Run(pattern);
      return 0;
    }
  }
  return RUN_ALL_TESTS();
}
#endif