added more unit tests

pdf_generator/generator.py

@@ -20,7 +20,7 @@ def copy_fact_image(target):
     shutil.copy(str(Path(__file__).parent / 'templates' / 'fact_logo.png'), str(Path(target) / 'fact_logo.png'))
 
 
-def generate_analysis_codes(engine, analysis):
+def generate_analysis_templates(engine, analysis):
     return [
         ('{}.tex'.format(analysis_plugin), engine.render_analysis_template(analysis_plugin, analysis[analysis_plugin])) for analysis_plugin in analysis
     ]
@@ -30,6 +30,21 @@ def create_report_filename(meta_data):
     return '{}_analysis_report.pdf'.format(meta_data['device_name'].replace(' ', '_').replace('/', '__'))
 
 
+def compile_pdf(meta_data, tmp_dir):
+    copy_fact_image(tmp_dir)
+    execute_latex(tmp_dir)
+    shutil.move(str(Path(tmp_dir, 'main.pdf')), str(Path(create_report_filename(meta_data))))
+
+
+def create_templates(analysis, meta_data, tmp_dir):
+    engine = Engine(tmp_dir=tmp_dir)
+
+    Path(tmp_dir, 'main.tex').write_text(engine.render_main_template(analysis=analysis, meta_data=meta_data))
+    Path(tmp_dir, 'meta.tex').write_text(engine.render_meta_template(meta_data))
+    for filename, result_code in generate_analysis_templates(engine=engine, analysis=analysis):
+        Path(tmp_dir, filename).write_text(result_code)
+
+
 def generate_report(firmware_uid, server_url=None):
     request_url = create_request_url(firmware_uid, server_url)
     try:
@@ -39,19 +54,7 @@ def generate_report(firmware_uid, server_url=None):
         return 1
 
     with TemporaryDirectory() as tmp_dir:
-        engine = Engine(tmp_dir=tmp_dir)
-        Path(tmp_dir, 'meta.tex').write_text(engine.render_meta_template(meta_data))
-
-        for filename, result_code in generate_analysis_codes(engine=engine, analysis=analysis):
-            Path(tmp_dir, filename).write_text(result_code)
-
-        Path(tmp_dir, 'main.tex').write_text(engine.render_main_template(analysis=analysis, meta_data=meta_data))
-
-        copy_fact_image(tmp_dir)
-
-        execute_latex(tmp_dir)
-
-        pdf_filename = create_report_filename(meta_data)
-        shutil.move(str(Path(tmp_dir, 'main.pdf')), str(Path('.', pdf_filename)))
+        create_templates(analysis, meta_data, tmp_dir)
+        compile_pdf(meta_data, tmp_dir)
 
     return 0

test/integration/test_code_generation.py

-import pytest
-from pdf_generator.tex_generation.template_engine import Engine
-from test.data.test_dict import TEST_DICT
-
-# pylint: disable=redefined-outer-name
-
-
-@pytest.fixture(scope='function')
-def stub_engine():
-    return Engine()
-
-
-def test_latex_code_generation(stub_engine: Engine):
-    result = stub_engine.render_meta_template(TEST_DICT)
-    assert result

test/unit/test_code_generation.py → test/integration/test_generator.py

-from tempfile import TemporaryDirectory
-
+import pytest
 from pdf_generator.generator import generate_report
 from pdf_generator.tex_generation.template_engine import Engine
+from test.data.test_dict import TEST_DICT
+
+# pylint: disable=redefined-outer-name
 
 TEST_DATA = {
     'analysis': {'file_hashes': {'ssdeep': 'bla', 'sha1': 'blah'}},
@@ -9,11 +11,20 @@ TEST_DATA = {
 }
 
 
-def test_render_template():
-    with TemporaryDirectory() as tmp_dir:
-        engine = Engine(template_folder='test', tmp_dir=tmp_dir)
-        test_data = {'meta_data': '123', 'analysis': '456'}
-        output = engine.render_analysis_template(plugin='render_test', analysis=test_data)
+@pytest.fixture(scope='function')
+def stub_engine():
+    return Engine()
+
+
+def test_latex_code_generation(stub_engine: Engine):
+    result = stub_engine.render_meta_template(TEST_DICT)
+    assert result
+
+
+def test_render_template(tmpdir):
+    engine = Engine(template_folder='test', tmp_dir=tmpdir)
+    test_data = {'meta_data': '123', 'analysis': '456'}
+    output = engine.render_analysis_template(plugin='render_test', analysis=test_data)
 
     assert output == 'Test  - '
 

test/unit/test_generator.py

+import json
+from pathlib import Path
+
+from pdf_generator.generator import (
+    copy_fact_image, create_report_filename, create_templates, execute_latex, generate_analysis_templates
+)
+
+
+class MockEngine:
+    def __init__(self, *_, **__):
+        pass
+
+    @staticmethod
+    def render_main_template(analysis, meta_data):
+        return '{}\n{}'.format(json.dumps(analysis), json.dumps(meta_data))
+
+    @staticmethod
+    def render_meta_template(meta_data):
+        return json.dumps(meta_data)
+
+    @staticmethod
+    def render_analysis_template(_, analysis):
+        return json.dumps(analysis)
+
+
+def exec_mock(*_, **__):
+    Path('test').write_text('works')
+
+
+def test_execute_latex(monkeypatch, tmpdir):
+    monkeypatch.setattr('pdf_generator.generator.execute_shell_command_get_return_code', exec_mock)
+
+    execute_latex(tmpdir)
+    assert Path(tmpdir, 'test').exists()
+    assert Path(tmpdir, 'test').read_text() == 'works'
+
+
+def test_copy_fact_image(tmpdir):
+    copy_fact_image(tmpdir)
+    assert Path(tmpdir, 'fact_logo.png').exists()
+
+
+def test_create_report_filename():
+    assert create_report_filename({'device_name': 'simple'}) == 'simple_analysis_report.pdf'
+    assert create_report_filename({'device_name': 'harder name'}) == 'harder_name_analysis_report.pdf'
+    assert create_report_filename({'device_name': 'dangerous/name'}) == 'dangerous__name_analysis_report.pdf'
+
+
+def test_create_analysis_templates():
+    templates = generate_analysis_templates(engine=MockEngine(), analysis={'test': {'result': 'data'}})
+
+    assert len(templates) == 1
+
+    filename, result_code = templates[0]
+    assert filename == 'test.tex'
+    assert result_code == '{"result": "data"}'
+
+
+def test_create_templates(monkeypatch, tmpdir):
+    monkeypatch.setattr('pdf_generator.generator.Engine', MockEngine)
+    create_templates(analysis={'test': {'result': 'data'}}, meta_data={}, tmp_dir=tmpdir)
+
+    assert Path(tmpdir, 'main.tex').exists()
+    assert Path(tmpdir, 'meta.tex').exists()
+    assert Path(tmpdir, 'test.tex').exists()
+
+    assert Path(tmpdir, 'test.tex').read_text() == '{"result": "data"}'