Changeset View
Changeset View
Standalone View
Standalone View
swh/graph/tests/test_cli.py
| # Copyright (C) 2019 The Software Heritage developers | # Copyright (C) 2019 The Software Heritage developers | ||||
| # See the AUTHORS file at the top-level directory of this distribution | # See the AUTHORS file at the top-level directory of this distribution | ||||
| # License: GNU General Public License version 3, or any later version | # License: GNU General Public License version 3, or any later version | ||||
| # See top-level LICENSE file for more information | # See top-level LICENSE file for more information | ||||
| import unittest | |||||
| from pathlib import Path | from pathlib import Path | ||||
| from tempfile import TemporaryDirectory, NamedTemporaryFile | from tempfile import TemporaryDirectory | ||||
| from typing import Dict | from typing import Dict | ||||
| from click.testing import CliRunner | from click.testing import CliRunner | ||||
| from swh.core import config | from swh.graph.cli import cli | ||||
| from swh.graph import cli | |||||
| DATA_DIR = Path(__file__).parents[0] / "dataset" | |||||
| def read_properties(properties_fname) -> Dict[str, str]: | def read_properties(properties_fname) -> Dict[str, str]: | ||||
| """read a Java .properties file""" | """read a Java .properties file""" | ||||
| properties = {} | |||||
| with open(properties_fname) as f: | with open(properties_fname) as f: | ||||
| for line in f: | keyvalues = ( | ||||
| if line.startswith("#"): | line.split("=", maxsplit=1) | ||||
| continue | for line in f | ||||
| (key, value) = line.rstrip().split("=", maxsplit=1) | if not line.strip().startswith("#") | ||||
| properties[key] = value | ) | ||||
| return dict((k.strip(), v.strip()) for (k, v) in keyvalues) | |||||
| return properties | |||||
| class TestCompress(unittest.TestCase): | |||||
| DATA_DIR = Path(__file__).parents[0] / "dataset" | |||||
| def setUp(self): | |||||
| self.runner = CliRunner() | |||||
| tmpconf = NamedTemporaryFile( | def test_pipeline(): | ||||
| mode="w", delete=False, prefix="swh-graph-test", suffix=".yml" | """run full compression pipeline""" | ||||
| ) | |||||
| # bare bone configuration, to allow testing the compression pipeline | # bare bone configuration, to allow testing the compression pipeline | ||||
| # with minimum RAM requirements on trivial graphs | # with minimum RAM requirements on trivial graphs | ||||
| tmpconf.write( | config = {"graph": {"compress": {"batch_size": 1000}}} | ||||
| """ | runner = CliRunner() | ||||
| graph: | |||||
| compress: | |||||
| batch_size: 1000 | |||||
| """ | |||||
| ) | |||||
| tmpconf.close() | |||||
| self.conffile = Path(tmpconf.name) | |||||
| self.config = config.read(self.conffile, cli.DEFAULT_CONFIG) | |||||
| def tearDown(self): | |||||
| if self.conffile.exists(): | |||||
| self.conffile.unlink() | |||||
| def test_pipeline(self): | |||||
| """run full compression pipeline""" | |||||
| with TemporaryDirectory(suffix=".swh-graph-test") as tmpdir: | with TemporaryDirectory(suffix=".swh-graph-test") as tmpdir: | ||||
| result = self.runner.invoke( | result = runner.invoke( | ||||
| cli.compress, | cli, | ||||
| ["--graph", self.DATA_DIR / "example", "--outdir", tmpdir], | ["compress", "--graph", DATA_DIR / "example", "--outdir", tmpdir], | ||||
| obj={"config": self.config}, | obj={"config": config}, | ||||
| ) | ) | ||||
| assert result.exit_code == 0, result | |||||
| self.assertEqual(result.exit_code, 0) | |||||
| properties = read_properties(Path(tmpdir) / "example.properties") | properties = read_properties(Path(tmpdir) / "example.properties") | ||||
| self.assertEqual(int(properties["nodes"]), 21) | |||||
| self.assertEqual(int(properties["arcs"]), 23) | assert int(properties["nodes"]) == 21 | ||||
| assert int(properties["arcs"]) == 23 | |||||