import os
import unittest

import yaml

from lib.db import DB
from lib.db_revision import DBRevision
from lib.tree_builder import TreeBuilder


class TestTreeMethods(unittest.TestCase):
    def setUp(self):
        self.db = DB(section="test")

    def verify_package(self, package):
        path = os.path.join(os.path.dirname(__file__), f"fixtures/{package}-data.yaml")
        with open(path, "r") as f:
            data = yaml.safe_load(f)
        for rev in data["revisions"]:
            DBRevision.import_fixture_dict(self.db, rev)

        revisions = TreeBuilder(self.db).build(package)
        path = os.path.join(
            os.path.dirname(__file__), f"fixtures/{package}-expected-tree.yaml"
        )
        # REGENERATE_DATA=1 PYTHONPATH=$PWD python3 ./tests/tree_test.py
        if os.getenv("REGENERATE_DATA"):
            with open(path, "w") as f:
                yaml.dump(revisions.as_list(), f)

        with open(path, "r") as f:
            data = yaml.safe_load(f)
        self.assertEqual(data, revisions.as_list())
        self.db.conn.rollback()

    def test_zsh_tree(self):
        self.verify_package("zsh")

    def test_clapper_tree(self):
        self.verify_package("clapper")


if __name__ == "__main__":
    unittest.main()