PyCS/test/test_database.py

import unittest
from contextlib import closing

from pycs import db
from pycs.database.Database import Database
from pycs.database.File import File
from pycs.database.Label import Label
from pycs.database.Model import Model
from pycs.database.LabelProvider import LabelProvider


class TestDatabase(unittest.TestCase):
    def setUp(self) -> None:
        db.create_all()
        # create database
        self.database = Database(discovery=False)

        # insert default models and label_providers
        with self.database:

            for i, supports in enumerate([["labeled-image", "fit"], ["labeled-bounding-boxes"], ["labeled-bounding-boxes"]], 1):

                model = Model.new(
                    name=f"Model {i}",
                    description=f"Description for Model {i}",
                    root_folder=f"modeldir{i}",
                )
                model.supports = supports

                if i > 2:
                    continue

                provider = LabelProvider.new(
                    name=f"Label Provider {i}",
                    description=f"Description for Label Provider {i}",
                    root_folder=f"labeldir{i}",
                )

        # projects
        models = list(self.database.models())
        label_providers = list(self.database.label_providers())

        for i, model in enumerate(models, 1):
            self.database.create_project(
                name=f'Project {i}',
                description=f'Project Description {i}',
                model=model,
                label_provider=label_providers[i-1] if i < 3 else None,
                root_folder=f'projectdir{i}',
                external_data=i==2,
                data_folder=f'datadir{i}',
            )

    def tearDown(self) -> None:
        db.drop_all()
        self.database.close()

    def test_models(self):
        models = list(self.database.models())

        # test length
        self.assertEqual(len(models), 3)

        # test insert
        for i in range(2):
            self.assertEqual(models[i].id, i + 1)
            self.assertEqual(models[i].name, f'Model {i + 1}')
            self.assertEqual(models[i].description, f'Description for Model {i + 1}')
            self.assertEqual(models[i].root_folder, f'modeldir{i + 1}')

        self.assertEqual(models[0].supports, ['labeled-image', 'fit'])
        self.assertEqual(models[1].supports, ['labeled-bounding-boxes'])

        # test copy
        copy, _ = models[0].copy_to('Copied Model', 'modeldir3')
        self.assertEqual(copy.id, 3)
        self.assertEqual(copy.name, 'Copied Model')
        self.assertEqual(copy.description, 'Description for Model 1')
        self.assertEqual(copy.root_folder, 'modeldir3')
        self.assertEqual(copy.supports, ['labeled-image', 'fit'])

    def test_label_providers(self):
        label_providers = list(self.database.label_providers())

        # test length
        self.assertEqual(len(label_providers), 2)

        for i in range(2):
            self.assertEqual(label_providers[i].id, i + 1)
            self.assertEqual(label_providers[i].name, f'Label Provider {i + 1}')
            self.assertEqual(label_providers[i].description, f'Description for Label Provider {i + 1}')
            self.assertEqual(label_providers[i].root_folder, f'labeldir{i + 1}')

    def test_projects(self):
        models = list(self.database.models())
        label_providers = list(self.database.label_providers())
        projects = list(self.database.projects())

        # create projects
        for i in range(3):
            project = projects[i]

            self.assertEqual(project.id, i + 1)
            self.assertEqual(project.name, f'Project {i + 1}')
            self.assertEqual(project.description, f'Project Description {i + 1}')
            self.assertEqual(project.model_id, i + 1)
            self.assertEqual(project.model.__dict__, models[i].__dict__)
            self.assertEqual(project.label_provider_id, label_providers[i].id if i < 2 else None)
            self.assertEqual(
                project.label_provider.__dict__ if project.label_provider is not None else None,
                label_providers[i].__dict__ if i < 2 else None
            )
            self.assertEqual(project.root_folder, f'projectdir{i + 1}')
            self.assertEqual(project.external_data, i == 1)
            self.assertEqual(project.data_folder, f'datadir{i + 1}')

        # get projects
        self.assertEqual(len(list(self.database.projects())), 3)

        # remove a project
        list(self.database.projects())[0].remove()
        projects = list(self.database.projects())

        self.assertEqual(len(projects), 2)
        self.assertEqual(projects[0].name, 'Project 2')

        # set properties
        project = list(self.database.projects())[0]

        project.set_name('Project 0')
        self.assertEqual(list(self.database.projects())[0].name, 'Project 0')

        project.set_description('Description 0')
        self.assertEqual(list(self.database.projects())[0].description, 'Description 0')


    def test_no_files_after_project_deletion(self):

        project = self.database.project(1)
        for i in range(5):
            file, is_new = project.add_file(
                uuid=f"some_string{i}",
                name=f"some_name{i}",
                filename=f"some_filename{i}",
                file_type="image",
                extension=".jpg",
                size=42,
            )

            self.assertTrue(is_new)
            self.assertIsNotNone(file)

        self.assertEqual(5, File.query.filter_by(project_id=project.id).count())

        project.remove()
        self.assertIsNone(self.database.project(1))
        self.assertEqual(0, File.query.filter_by(project_id=project.id).count())

    def test_no_labels_after_project_deletion(self):

        project = self.database.project(1)
        for i in range(5):
            label, is_new = project.create_label(
                name=f"label{i}",
                reference=f"ref{i}"
            )

            self.assertTrue(is_new)
            self.assertIsNotNone(label)

        self.assertEqual(5, Label.query.filter_by(project_id=project.id).count())

        project.remove()
        self.assertIsNone(self.database.project(1))
        self.assertEqual(0, Label.query.filter_by(project_id=project.id).count())

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