PyCS/tests/base.py

import os
import shutil
import unittest
import typing as T

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

server = None

class BaseTestCase(unittest.TestCase):

    def setUp(self, discovery: bool = False):
        global server
        app.config["TESTING"] = True
        self.projects_dir = app.config["TEST_PROJECTS_DIR"] = "test_projects"
        app.config["WTF_CSRF_ENABLED"] = False
        app.config["DEBUG"] = False
        app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///"

        db.create_all()

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

        # init the server once
        if server is None:
            server = WebServer(app, settings, discovery)

        elif discovery:
            # run discovery modules manually
            Model.discover("models/")
            LabelProvider.discover("labels/")

        self.setupModels()

    def setupModels(self):
        pass

    def tearDown(self):
        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)
        self.assertEqual(response.status_code, status_code, response.get_data().decode())
        return response

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

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

    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,
            content_type=content_type,
        )