import pathlib

import subprocess

import sys

import pytest

from misc.codegen.generators import qlgen

from misc.codegen.lib import ql

from misc.codegen.test.utils import *

@pytest.fixture(autouse=True)

def run_mock():

with mock.patch("subprocess.run") as ret:

ret.return_value.returncode = 0

yield ret

# these are lambdas so that they will use patched paths when called

def stub_path():

return paths.root_dir / "ql/lib/stub/path"

def ql_output_path():

return paths.root_dir / "ql/lib/other/path"

def ql_test_output_path():

return paths.root_dir / "ql/test/path"

def generated_registry_path():

return paths.root_dir / "ql/registry.list"

def import_file():

return stub_path().with_suffix(".qll")

def children_file():

return ql_output_path() / "ParentChild.qll"

stub_import = "stub.path"

stub_import_prefix_internal = stub_import + ".internal."

stub_import_prefix = stub_import + "."

root_import = stub_import_prefix_internal + "Element"

gen_import = "other.path"

gen_import_prefix = gen_import + "."

@pytest.fixture

def qlgen_opts(opts):

opts.ql_stub_output = stub_path()

opts.ql_output = ql_output_path()

opts.ql_test_output = ql_test_output_path()

opts.generated_registry = generated_registry_path()

opts.ql_format = True

opts.root_dir = paths.root_dir

opts.force = False

opts.codeql_binary = "./my_fake_codeql"

pathlib.Path(opts.codeql_binary).touch()

return opts

@pytest.fixture

def generate(input, qlgen_opts, renderer, render_manager):

render_manager.written = []

def func(classes):

input.classes = {cls.name: cls for cls in classes}

return run_managed_generation(

qlgen.generate, qlgen_opts, renderer, render_manager

)

return func

@pytest.fixture

def generate_import_list(generate):

def func(classes):

ret = generate(classes)

assert import_file() in ret

return ret[import_file()]

return func

@pytest.fixture

def generate_children_implementations(generate):

def func(classes):

ret = generate(classes)

assert children_file() in ret

return ret[children_file()]

return func

def _filter_generated_classes(ret, output_test_files=False):

files = {x for x in ret}

print(files)

files.remove(import_file())

files.remove(children_file())

stub_files = set()

base_files = set()

test_files = set()

for f in files:

try:

stub_files.add(f.relative_to(stub_path()))

print(f)

except ValueError:

try:

base_files.add(f.relative_to(ql_output_path()))

except ValueError:

try:

test_files.add(f.relative_to(ql_test_output_path()))

except ValueError:

assert False, f"{f} is in wrong directory"

if output_test_files:

return {str(f): ret[ql_test_output_path() / f] for f in test_files}

base_files -= {

pathlib.Path(f"{name}.qll")

for name in ("Raw", "Synth", "SynthConstructors", "PureSynthConstructors")

}

stub_files = {

(

pathlib.Path(f.parent.parent, f.stem + ".qll")

if f.parent.name == "internal"

and pathlib.Path(f.parent.parent, f.stem + ".qll") in base_files

else f

)

for f in stub_files

}

assert base_files <= stub_files

return {

str(f): (

(

ret[stub_path() / "internal" / f]

if stub_path() / "internal" / f in ret

else ret[stub_path() / f]

),

ret[

stub_path()

/ pathlib.Path(

f.parent,

"internal" if not f.parent.name == "internal" else "",

f.stem + "Impl.qll",

)

],

ret[ql_output_path() / f],

)

for f in base_files

}

@pytest.fixture

def generate_classes(generate):

def func(classes):

return _filter_generated_classes(generate(classes))

return func

@pytest.fixture

def generate_tests(generate):

def func(classes):

return _filter_generated_classes(generate(classes), output_test_files=True)

return func

def a_ql_class(**kwargs):

return ql.Class(**kwargs, import_prefix=gen_import)

def a_ql_stub(*, name, import_prefix="", **kwargs):

return ql.Stub(

name=name,

**kwargs,

import_prefix=gen_import,

base_import=f"{gen_import_prefix}{import_prefix}{name}",

)

def a_ql_class_public(*, name, **kwargs):

return ql.ClassPublic(name=name, **kwargs)

def test_one_empty_class(generate_classes):

assert generate_classes([schema.Class("A")]) == {

"A.qll": (

a_ql_class_public(name="A"),

a_ql_stub(name="A"),

a_ql_class(name="A", final=True, imports=[stub_import_prefix + "A"]),

)

}

def test_one_empty_internal_class(generate_classes):

assert generate_classes([schema.Class("A", pragmas=["ql_internal"])]) == {

"A.qll": (

a_ql_class_public(name="A", internal=True),

a_ql_stub(name="A"),

a_ql_class(

name="A",

final=True,

internal=True,

imports=[stub_import_prefix_internal + "A"],

),

),

}

def test_hierarchy(generate_classes):

assert generate_classes(

[

schema.Class("D", bases=["B", "C"]),

schema.Class("C", bases=["A"], derived={"D"}),

schema.Class("B", bases=["A"], derived={"D"}),

schema.Class("A", derived={"B", "C"}),

]

) == {

"A.qll": (

a_ql_class_public(name="A"),

a_ql_stub(name="A"),

a_ql_class(name="A", imports=[stub_import_prefix + "A"]),

),

"B.qll": (

a_ql_class_public(name="B", imports=[stub_import_prefix + "A"]),

a_ql_stub(name="B"),

a_ql_class(

name="B",

bases=["A"],

bases_impl=["AImpl::A"],

imports=[stub_import_prefix_internal + "AImpl::Impl as AImpl"],

),

),

"C.qll": (

a_ql_class_public(name="C", imports=[stub_import_prefix + "A"]),

a_ql_stub(name="C"),

a_ql_class(

name="C",

bases=["A"],

bases_impl=["AImpl::A"],

imports=[stub_import_prefix_internal + "AImpl::Impl as AImpl"],

),

),

"D.qll": (

a_ql_class_public(

name="D", imports=[stub_import_prefix + "B", stub_import_prefix + "C"]

),

a_ql_stub(name="D"),

a_ql_class(

name="D",

final=True,

bases=["B", "C"],

bases_impl=["BImpl::B", "CImpl::C"],

imports=[

stub_import_prefix_internal + cls + "Impl::Impl as " + cls + "Impl"

for cls in "BC"

],

),

),

}

def test_hierarchy_imports(generate_import_list):

assert generate_import_list(

[

schema.Class("D", bases=["B", "C"]),

schema.Class("C", bases=["A"], derived={"D"}),

schema.Class("B", bases=["A"], derived={"D"}),

schema.Class("A", derived={"B", "C"}),

]

) == ql.ImportList([stub_import_prefix + cls for cls in "ABCD"])

def test_internal_not_in_import_list(generate_import_list):

assert generate_import_list(

[

schema.Class("D", bases=["B", "C"]),

schema.Class("C", bases=["A"], derived={"D"}, pragmas=["ql_internal"]),

schema.Class("B", bases=["A"], derived={"D"}),

schema.Class("A", derived={"B", "C"}, pragmas=["ql_internal"]),

]

) == ql.ImportList([stub_import_prefix + cls for cls in "BD"])

def test_hierarchy_children(generate_children_implementations):

assert generate_children_implementations(

[

schema.Class("A", derived={"B", "C"}, pragmas=["ql_internal"]),

schema.Class("B", bases=["A"], derived={"D"}),

schema.Class("C", bases=["A"], derived={"D"}, pragmas=["ql_internal"]),

schema.Class("D", bases=["B", "C"]),

]

) == ql.GetParentImplementation(

classes=[

a_ql_class(

name="A", internal=True, imports=[stub_import_prefix_internal + "A"]

),

a_ql_class(

name="B",

bases=["A"],

bases_impl=["AImpl::A"],

imports=[stub_import_prefix_internal + "AImpl::Impl as AImpl"],

),

a_ql_class(

name="C",

bases=["A"],

bases_impl=["AImpl::A"],

imports=[stub_import_prefix_internal + "AImpl::Impl as AImpl"],

internal=True,

),

a_ql_class(

name="D",

final=True,

bases=["B", "C"],

bases_impl=["BImpl::B", "CImpl::C"],

imports=[

stub_import_prefix_internal + cls + "Impl::Impl as " + cls + "Impl"

for cls in "BC"

],

),

],

imports=[stub_import] + [stub_import_prefix_internal + cls for cls in "AC"],

)

def test_single_property(generate_classes):

assert generate_classes(

[

schema.Class("MyObject", properties=[schema.SingleProperty("foo", "bar")]),

]

) == {

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="Foo",

type="bar",

tablename="my_objects",

tableparams=["this", "result"],

doc="foo of this my object",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_internal_property(generate_classes):

assert generate_classes(

[

schema.Class(

"MyObject",

properties=[

schema.SingleProperty("foo", "bar", pragmas=["ql_internal"])

],

),

]

) == {

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="Foo",

type="bar",

tablename="my_objects",

tableparams=["this", "result"],

doc="foo of this my object",

internal=True,

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_children(generate_classes):

expected_parent_property = ql.Property(

singular="ParentChild",

type="int",

is_child=True,

tablename="parents",

tableparams=["this", "result"],

doc="parent child of this parent",

)

expected_properties = [

ql.Property(

singular="A",

type="int",

tablename="my_objects",

tableparams=["this", "result", "_"],

doc="a of this my object",

),

ql.Property(

singular="Child1",

type="int",

tablename="my_objects",

tableparams=["this", "_", "result"],

is_child=True,

doc="child 1 of this my object",

),

ql.Property(

singular="B",

plural="Bs",

type="int",

tablename="my_object_bs",

tableparams=["this", "index", "result"],

doc="b of this my object",

doc_plural="bs of this my object",

),

ql.Property(

singular="Child",

plural="Children",

type="int",

tablename="my_object_children",

tableparams=["this", "index", "result"],

is_child=True,

doc="child of this my object",

doc_plural="children of this my object",

),

ql.Property(

singular="C",

type="int",

tablename="my_object_cs",

tableparams=["this", "result"],

is_optional=True,

doc="c of this my object",

),

ql.Property(

singular="Child3",

type="int",

tablename="my_object_child_3s",

tableparams=["this", "result"],

is_optional=True,

is_child=True,

doc="child 3 of this my object",

),

ql.Property(

singular="D",

plural="Ds",

type="int",

tablename="my_object_ds",

tableparams=["this", "index", "result"],

is_optional=True,

doc="d of this my object",

doc_plural="ds of this my object",

),

ql.Property(

singular="Child4",

plural="Child4s",

type="int",

tablename="my_object_child_4s",

tableparams=["this", "index", "result"],

is_optional=True,

is_child=True,

doc="child 4 of this my object",

doc_plural="child 4s of this my object",

),

]

assert generate_classes(

[

schema.Class("FakeRoot"),

schema.Class(

"Parent",

derived={"MyObject"},

properties=[

schema.SingleProperty("parent_child", "int", is_child=True),

],

),

schema.Class(

"MyObject",

bases=["Parent"],

properties=[

schema.SingleProperty("a", "int"),

schema.SingleProperty("child_1", "int", is_child=True),

schema.RepeatedProperty("bs", "int"),

schema.RepeatedProperty("children", "int", is_child=True),

schema.OptionalProperty("c", "int"),

schema.OptionalProperty("child_3", "int", is_child=True),

schema.RepeatedOptionalProperty("d", "int"),

schema.RepeatedOptionalProperty("child_4", "int", is_child=True),

],

),

]

) == {

"FakeRoot.qll": (

a_ql_class_public(name="FakeRoot"),

a_ql_stub(name="FakeRoot"),

a_ql_class(

name="FakeRoot", final=True, imports=[stub_import_prefix + "FakeRoot"]

),

),

"Parent.qll": (

a_ql_class_public(name="Parent"),

a_ql_stub(name="Parent"),

a_ql_class(

name="Parent",

imports=[stub_import_prefix + "Parent"],

properties=[expected_parent_property],

all_children=[

ql.Child(

expected_parent_property,

),

],

),

),

"MyObject.qll": (

a_ql_class_public(name="MyObject", imports=[stub_import_prefix + "Parent"]),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

bases=["Parent"],

bases_impl=["ParentImpl::Parent"],

properties=expected_properties,

all_children=[

ql.Child(

expected_parent_property,

),

ql.Child(

expected_properties[1],

prev="ParentChild",

),

ql.Child(

expected_properties[3],

prev="Child1",

),

ql.Child(

expected_properties[5],

prev="Child",

),

ql.Child(

expected_properties[7],

prev="Child3",

),

],

imports=[

stub_import_prefix + "internal.ParentImpl::Impl as ParentImpl"

],

),

),

}

def test_single_properties(generate_classes):

assert generate_classes(

[

schema.Class(

"MyObject",

properties=[

schema.SingleProperty("one", "x"),

schema.SingleProperty("two", "y"),

schema.SingleProperty("three", "z"),

],

),

]

) == {

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="One",

type="x",

tablename="my_objects",

tableparams=["this", "result", "_", "_"],

doc="one of this my object",

),

ql.Property(

singular="Two",

type="y",

tablename="my_objects",

tableparams=["this", "_", "result", "_"],

doc="two of this my object",

),

ql.Property(

singular="Three",

type="z",

tablename="my_objects",

tableparams=["this", "_", "_", "result"],

doc="three of this my object",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_optional_property(generate_classes):

assert generate_classes(

[

schema.Class("FakeRoot"),

schema.Class(

"MyObject",

properties=[schema.OptionalProperty("foo", "bar")],

),

]

) == {

"FakeRoot.qll": (

a_ql_class_public(name="FakeRoot"),

a_ql_stub(name="FakeRoot"),

a_ql_class(

name="FakeRoot", final=True, imports=[stub_import_prefix + "FakeRoot"]

),

),

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="Foo",

type="bar",

tablename="my_object_foos",

tableparams=["this", "result"],

is_optional=True,

doc="foo of this my object",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_repeated_property(generate_classes):

assert generate_classes(

[

schema.Class("FakeRoot"),

schema.Class(

"MyObject",

properties=[schema.RepeatedProperty("foo", "bar")],

),

]

) == {

"FakeRoot.qll": (

a_ql_class_public(name="FakeRoot"),

a_ql_stub(name="FakeRoot"),

a_ql_class(

name="FakeRoot", final=True, imports=[stub_import_prefix + "FakeRoot"]

),

),

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="Foo",

plural="Foos",

type="bar",

tablename="my_object_foos",

tableparams=["this", "index", "result"],

doc="foo of this my object",

doc_plural="foos of this my object",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_repeated_unordered_property(generate_classes):

assert generate_classes(

[

schema.Class("FakeRoot"),

schema.Class(

"MyObject", properties=[schema.RepeatedUnorderedProperty("foo", "bar")]

),

]

) == {

"FakeRoot.qll": (

a_ql_class_public(name="FakeRoot"),

a_ql_stub(name="FakeRoot"),

a_ql_class(

name="FakeRoot", final=True, imports=[stub_import_prefix + "FakeRoot"]

),

),

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="Foo",

plural="Foos",

type="bar",

tablename="my_object_foos",

tableparams=["this", "result"],

is_unordered=True,

doc="foo of this my object",

doc_plural="foos of this my object",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_repeated_optional_property(generate_classes):

assert generate_classes(

[

schema.Class("FakeRoot"),

schema.Class(

"MyObject",

properties=[schema.RepeatedOptionalProperty("foo", "bar")],

),

]

) == {

"FakeRoot.qll": (

a_ql_class_public(name="FakeRoot"),

a_ql_stub(name="FakeRoot"),

a_ql_class(

name="FakeRoot", final=True, imports=[stub_import_prefix + "FakeRoot"]

),

),

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="Foo",

plural="Foos",

type="bar",

tablename="my_object_foos",

tableparams=["this", "index", "result"],

is_optional=True,

doc="foo of this my object",

doc_plural="foos of this my object",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_predicate_property(generate_classes):

assert generate_classes(

[

schema.Class("MyObject", properties=[schema.PredicateProperty("is_foo")]),

]

) == {

"MyObject.qll": (

a_ql_class_public(name="MyObject"),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

properties=[

ql.Property(

singular="isFoo",

type="predicate",

tablename="my_object_is_foo",

tableparams=["this"],

is_predicate=True,

doc="this my object is foo",

),

],

imports=[stub_import_prefix + "MyObject"],

),

),

}

def test_single_class_property(generate_classes):

assert generate_classes(

[

schema.Class("Bar"),

schema.Class(

"MyObject",

properties=[schema.SingleProperty("foo", "Bar")],

),

]

) == {

"MyObject.qll": (

a_ql_class_public(name="MyObject", imports=[stub_import_prefix + "Bar"]),

a_ql_stub(name="MyObject"),

a_ql_class(

name="MyObject",

final=True,

imports=[stub_import_prefix + "Bar", stub_import_prefix + "MyObject"],

properties=[

ql.Property(

singular="Foo",

type="Bar",

tablename="my_objects",

tableparams=["this", "result"],

doc="foo of this my object",

type_is_codegen_class=True,

),

],

),

),

"Bar.qll": (

a_ql_class_public(name="Bar"),

a_ql_stub(name="Bar"),

a_ql_class(name="Bar", final=True, imports=[stub_import_prefix + "Bar"]),

),

}

def test_class_with_doc(generate_classes):

doc = ["Very important class.", "Very."]

assert generate_classes(

[

schema.Class("A", doc=doc),

]

) == {

"A.qll": (

a_ql_class_public(name="A", doc=doc),

a_ql_stub(name="A", doc=doc),

a_ql_class(

name="A", final=True, doc=doc, imports=[stub_import_prefix + "A"]

),

),

}

def test_class_dir(generate_classes):

dir = "another/rel/path"

assert generate_classes(

[

schema.Class("A", derived={"B"}, pragmas={"group": dir}),

schema.Class("B", bases=["A"]),

]

) == {

f"{dir}/A.qll": (

a_ql_class_public(name="A"),

a_ql_stub(name="A", import_prefix="another.rel.path."),

a_ql_class(

name="A",

dir=pathlib.Path(dir),

imports=[stub_import_prefix + "another.rel.path.A"],

),

),

"B.qll": (

a_ql_class_public(

name="B", imports=[stub_import_prefix + "another.rel.path.A"]

),

a_ql_stub(name="B"),

a_ql_class(

name="B",

final=True,

bases=["A"],

bases_impl=["AImpl::A"],

imports=[

stub_import_prefix

+ "another.rel.path.internal.AImpl::Impl as AImpl"

],

),

),

}

def test_root_element_cannot_have_children(generate_classes):

with pytest.raises(qlgen.RootElementHasChildren):

generate_classes(

[schema.Class("A", properties=[schema.SingleProperty("x", is_child=True)])]

)

def test_class_dir_imports(generate_import_list):

dir = "another/rel/path"

assert generate_import_list(

[

schema.Class("A", derived={"B"}, pragmas={"group": dir}),

schema.Class("B", bases=["A"]),

]

) == ql.ImportList(

[

stub_import_prefix + "B",

stub_import_prefix + "another.rel.path.A",

]

)

def test_format(opts, generate, render_manager, run_mock):

run_mock.return_value.stderr = "some\nlines\n"

render_manager.written = [

pathlib.Path("x", "foo.ql"),

pathlib.Path("bar.qll"),

pathlib.Path("y", "baz.txt"),

]

generate([schema.Class("A")])

assert run_mock.mock_calls == [

mock.call(

[

opts.codeql_binary,

"query",

"format",

"--in-place",

"--",

"x/foo.ql",

"bar.qll",

],

stderr=subprocess.PIPE,

text=True,

),

]

def test_format_error(opts, generate, render_manager, run_mock):

run_mock.return_value.stderr = "some\nlines\n"

run_mock.return_value.returncode = 1

render_manager.written = [

pathlib.Path("x", "foo.ql"),

pathlib.Path("bar.qll"),

pathlib.Path("y", "baz.txt"),

]

with pytest.raises(qlgen.FormatError):

generate([schema.Class("A")])

def test_format_no_codeql(opts, generate, render_manager, run_mock):

pathlib.Path(opts.codeql_binary).unlink()

render_manager.written = [

pathlib.Path("bar.qll"),

]

with pytest.raises(qlgen.FormatError):

generate([schema.Class("A")])

def test_format_no_codeql_in_path(opts, generate, render_manager, run_mock):

opts.codeql_binary = "my_fake_codeql"

render_manager.written = [

pathlib.Path("bar.qll"),

]

with pytest.raises(qlgen.FormatError):

generate([schema.Class("A")])

@pytest.mark.parametrize("force", [False, True])

def test_manage_parameters(opts, generate, renderer, force):

opts.force = force

ql_a = opts.ql_output / "A.qll"

ql_b = opts.ql_output / "B.qll"

stub_a = opts.ql_stub_output / "A.qll"

stub_b = opts.ql_stub_output / "B.qll"

test_a = opts.ql_test_output / "A.ql"

test_b = opts.ql_test_output / "MISSING_SOURCE.txt"

test_c = opts.ql_test_output / "B.txt"

write(ql_a)

write(ql_b)

write(stub_a)

write(stub_b)

write(test_a)

write(test_b)

write(test_c)

generate([schema.Class("A")])

assert renderer.mock_calls == [

mock.call.manage(

generated={ql_a, ql_b, test_a, test_b, import_file()},

stubs={stub_a, stub_b},

registry=opts.generated_registry,

force=force,

)

]

def test_modified_stub_skipped(qlgen_opts, generate, render_manager):

stub = qlgen_opts.ql_stub_output / "AImpl.qll"

render_manager.is_customized_stub.side_effect = lambda f: f == stub

assert stub not in generate([schema.Class("A")])

def test_test_missing_source(generate_tests):

generate_tests(

[

schema.Class("A"),

]

) == {