PyCS/tests/base.py

import os
import shutil
import unittest
import eventlet
import typing as T

from unittest import mock

from pycs import app
from pycs import db
from pycs import settings
from pycs.frontend.WebServer import WebServer
from pycs.database.Model import Model
from pycs.database.LabelProvider import LabelProvider
from pycs.util.PipelineCache import PipelineCache

server = None

def pаtch_tpool_execute(test_func):

    def call_func(func, *args, **kwargs):
        return func(*args, **kwargs)

    decorator = mock.patch("eventlet.tpool.execute",
        side_effect=call_func)

    return decorator(test_func)

class BaseTestCase(unittest.TestCase):
    _sleep_time = 0.2
    server = None

    @classmethod
    def setUpClass(cls, discovery: bool = False):
        global server
        PipelineCache.CLOSE_TIMER = 2
        app.config["TESTING"] = True
        app.config["WTF_CSRF_ENABLED"] = False
        app.config["DEBUG"] = False
        app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///"

        if server is None:
            server = WebServer(app, settings, discovery)

        if cls.server is None:
            cls.server = server
        db.create_all()

        # # run discovery modules manually
        # Model.discover("models/")
        # LabelProvider.discover("labels/")
        cls.server.pipelines.start()


    def wait_for_bg_jobs(self):

        # wait for JobRunner jobs to finish
        while True:
            ready = True
            for job in self.server.jobs.list():
                if job.finished is None:
                    print(f"{job} is not finished!")
                    ready = False
                    break

            if ready:
                break

            eventlet.sleep(self._sleep_time)

        # wait for PipelineCache to finish

        while not self.server.pipelines.is_empty:
            eventlet.sleep(self._sleep_time)

    @classmethod
    def tearDownClass(cls):
        super().tearDownClass()
        cls.server.pipelines.shutdown()



    def setUp(self):
        self.projects_dir = app.config["TEST_PROJECTS_DIR"] = "test_projects"

        db.create_all()

        self.client = app.test_client()
        self.context = app.test_request_context()
        self.context.push()

        self.setupModels()

    def setupModels(self):
        pass

    def tearDown(self):
        self.wait_for_bg_jobs()

        self.context.pop()

        if os.path.exists(self.projects_dir):
            shutil.rmtree(self.projects_dir)

        db.drop_all()


    def _do_request(self, request_func, *args,  status_code=200, **kwargs):
        response = request_func(*args, follow_redirects=True, **kwargs)
        try:
            content = response.get_data().decode()
        except UnicodeDecodeError:
            content = response.get_data()

        self.assertEqual(response.status_code, status_code, content)
        return response

    def post(self,
             url: str,
             *,
             status_code: int = 200,
             data: T.Optional[dict] = None,
             json: T.Optional[dict] = None,
             **kwargs):

        return self._do_request(
            self.client.post,
            url,
            status_code=status_code,
            json=json,
            data=data,
            **kwargs
        )

    def get(self,
             url: str,
             *,
             status_code: int = 200,
             data: T.Optional[dict] = None,
             json: T.Optional[dict] = None):

        return self._do_request(
            self.client.get,
            url,
            status_code=status_code,
            json=json,
            data=data,
        )

    def wait_for_coroutines(self):
        eventlet.sleep(self._sleep_time)