xref: /llvm-project/llvm/utils/mlgo-utils/tests/corpus/make_corpus_test.py (revision d3ca779d7264bf0751eb4db632c14d8bc05608ca) 
1# REQUIRES: system-linux
2
3## Test the functionality of make_corpus_lib
4
5import json
6import os
7import sys
8
9from mlgo.corpus import make_corpus_lib
10
11## Test that when we load the bitcode from a directory using the
12## load_bitcode_from_directory function, we get the expected results.
13
14# RUN: rm -rf %t.dir && mkdir %t.dir
15# RUN: mkdir %t.dir/nested
16# RUN: touch %t.dir/nested/test1.bc
17# RUN: touch %t.dir/nested/test2.bc
18# RUN: %python %s test_load_bitcode_from_directory %t.dir | FileCheck %s --check-prefix CHECK-LOAD
19
20
21def test_load_bitcode_from_directory(work_dir):
22    relative_paths = make_corpus_lib.load_bitcode_from_directory(work_dir)
23    relative_paths = sorted(relative_paths)
24    for relative_path in relative_paths:
25        print(relative_path)
26    # CHECK-LOAD: nested/test1
27    # CHECK-LOAD: nested/test2
28
29
30## Test that when we copy the bitcode given a list of relative paths, the
31## appropriate files are copied over.
32
33# RUN: rm -rf %t.dir1 && mkdir %t.dir1
34# RUN: %python %s test_copy_bitcode %t.dir %t.dir1
35# RUN: ls %t.dir1/nested | FileCheck %s --check-prefix CHECK-COPY
36
37# CHECK-COPY: test1.bc
38# CHECK-COPY: test2.bc
39
40
41def test_copy_bitcode(directory, output_dir):
42    relative_paths = ["nested/test1", "nested/test2"]
43    make_corpus_lib.copy_bitcode(relative_paths, directory, output_dir)
44
45
46## Test that we get the expected corpus manifest when writing a corpus
47## manifest to the specificed directory.
48
49# RUN: %python %s test_write_corpus_manifest %t.dir1 | FileCheck %s --check-prefix CHECK-MANIFEST
50
51
52def test_write_corpus_manifest(output_dir):
53    relative_output_paths = ["test/test1", "test/test2"]
54    default_args = ["-O3", "-c"]
55    make_corpus_lib.write_corpus_manifest(
56        relative_output_paths, output_dir, default_args
57    )
58    with open(
59        os.path.join(output_dir, "corpus_description.json"), encoding="utf-8"
60    ) as corpus_description_file:
61        corpus_description = json.load(corpus_description_file)
62    print(corpus_description["global_command_override"])
63    # CHECK-MANIFEST: ['-O3', '-c']
64    print(corpus_description["has_thinlto"])
65    # CHECK-MANIFEST: False
66    print(corpus_description["modules"])
67    # CHECK-MANIFEST: ['test/test1', 'test/test2']
68
69
70if __name__ == "__main__":
71    globals()[sys.argv[1]](*sys.argv[2:])
72