conftest.py

import json
import os
import tempfile
import random

import pytest
import compendium_v2
from compendium_v2 import db
from compendium_v2.db import model
from compendium_v2 import survey_db
from compendium_v2.survey_db import model as survey_model

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool

import csv


def _test_data_csv(filename):
    data_filename = os.path.join(
        os.path.dirname(__file__),
        'data',
        filename)
    yield from csv.DictReader(open(data_filename, "r"))


@pytest.fixture
def dummy_config():
    yield {
        'SQLALCHEMY_DATABASE_URI': 'sqlite://',
        'SURVEY_DATABASE_URI': 'sqlite:///'
    }


@pytest.fixture
def mocked_survey_db(mocker):
    engine = create_engine(
        'sqlite://',
        connect_args={'check_same_thread': False},
        poolclass=StaticPool,
        echo=False)
    survey_model.base_schema.metadata.create_all(engine)
    mocker.patch(
        'compendium_v2.survey_db._SESSION_MAKER',
        sessionmaker(bind=engine))
    mocker.patch(
        'compendium_v2.survey_db.init_db_model',
        lambda dsn: None)


@pytest.fixture
def mocked_db(mocker):
    # cf. https://stackoverflow.com/a/33057675
    engine = create_engine(
        'sqlite://',
        connect_args={'check_same_thread': False},
        poolclass=StaticPool,
        echo=False)
    model.base_schema.metadata.create_all(engine)
    mocker.patch(
        'compendium_v2.db._SESSION_MAKER',
        sessionmaker(bind=engine))
    mocker.patch(
        'compendium_v2.db.init_db_model',
        lambda dsn: None)
    mocker.patch(
        'compendium_v2.migrate_database',
        lambda config: None)


@pytest.fixture
def test_budget_data():
    with db.session_scope() as session:
        data = [row for row in _test_data_csv("BudgetTestData.csv")]
        nren_names = set([row["nren"] for row in data])
        nren_dict = {
            nren_name: model.NREN(name=nren_name) for nren_name in nren_names}
        session.add_all(nren_dict.values())

        for row in data:
            nren = nren_dict[row["nren"]]
            budget = row["budget"]
            year = row["year"]

            session.add(
                model.BudgetEntry(
                    nren=nren,
                    budget=float(budget),
                    year=int(year))
            )
    with survey_db.session_scope() as session:
        data = _test_data_csv("BudgetTestData.csv")
        nrens = set()
        budgets_data = []
        for row in data:
            nren = row["nren"]
            budget = row["budget"]
            year = row["year"]
            country_code = row["nren"]

            nrens.add(nren)

            budgets_data.append(
                survey_model.Budgets(
                    budget=budget,
                    year=year,
                    country_code=country_code
                )
            )
        for nren in nrens:
            session.add(
                survey_model.Nrens(
                    abbreviation=nren,
                    country_code=nren
                ))

        session.add_all(budgets_data)


@pytest.fixture
def test_funding_source_data():
    with db.session_scope() as session:
        data = [row for row in _test_data_csv("FundingSourceTestData.csv")]
        nren_names = set([row["nren"] for row in data])
        nren_dict = {
            nren_name: model.NREN(name=nren_name) for nren_name in nren_names}
        session.add_all(nren_dict.values())

        for row in data:
            nren = nren_dict[row["nren"]]
            year = row["year"]
            client = row["client"]
            european = row["european"]
            publicbodies = row["publicbodies"]
            commercial = row["commercial"]
            other = row["other"]

            session.add(
                model.FundingSource(
                    nren=nren, year=year,
                    client_institutions=client,
                    european_funding=european,
                    gov_public_bodies=publicbodies,
                    commercial=commercial,
                    other=other)
            )


@pytest.fixture
def test_staff_data():

    # generator of random test data for 5 years and 100 nrens
    def _generate_rows():
        for nren in ["nren" + str(i) for i in range(1, 100)]:
            for year in range(2016, 2021):
                yield {
                    "nren": nren,
                    "year": year,
                    "permanent_fte": random.randint(0, 100),
                    "subcontracted_fte": random.randint(0, 100),
                    "technical_fte": random.randint(0, 100),
                    "non_technical_fte": random.randint(0, 100)
                }

    with db.session_scope() as session:

        data = list(_generate_rows())

        nren_dict = {
            nren_name: model.NREN(name=nren_name) for nren_name in [d['nren'] for d in data]}
        session.add_all(nren_dict.values())

        for row in data:
            nren = nren_dict[row["nren"]]
            year = row["year"]
            permanent_fte = row["permanent_fte"]
            subcontracted_fte = row["subcontracted_fte"]
            technical_fte = row["technical_fte"]
            non_technical_fte = row["non_technical_fte"]

            session.add(
                model.NrenStaff(
                    nren=nren,
                    year=year,
                    permanent_fte=permanent_fte,
                    subcontracted_fte=subcontracted_fte,
                    technical_fte=technical_fte,
                    non_technical_fte=non_technical_fte
                )
            )


@pytest.fixture
def data_config_filename(dummy_config):
    with tempfile.NamedTemporaryFile() as f:
        f.write(json.dumps(dummy_config).encode('utf-8'))
        f.flush()
        yield f.name


@pytest.fixture
def client(data_config_filename, mocked_db, mocked_survey_db):
    os.environ['SETTINGS_FILENAME'] = data_config_filename
    with compendium_v2.create_app().test_client() as c:
        yield c


@pytest.fixture
def test_charging_structure_data():
    with db.session_scope() as session:
        data = [row for row in _test_data_csv("ChargingStructureTestData.csv")]
        nren_names = set([row["nren"] for row in data])
        nren_dict = {
            nren_name: model.NREN(name=nren_name) for nren_name in nren_names}
        session.add_all(nren_dict.values())

        for row in data:
            nren = nren_dict[row["nren"]]
            year = row["year"]
            fee_type = row["fee_type"]
            if fee_type == "null":
                fee_type = None

            session.add(
                model.ChargingStructure(
                    nren=nren, year=year,
                    fee_type=fee_type)
            )


@pytest.fixture
def test_organization_data():
    def _generate_sub_org_data():
        for nren in ["nren" + str(i) for i in range(1, 50)]:
            for year in range(2016, 2021):
                yield {
                    'nren': nren,
                    'year': year,
                    'name': 'sub_org' + str(random.randint(1, 100)),
                    'role': random.choice(['technical centre', 'stuff', 'test123']),
                }

    def _generate_org_data():
        for nren in ["nren" + str(i) for i in range(1, 50)]:
            for year in range(2016, 2021):
                yield {
                    'nren': nren,
                    'year': year,
                    'name': 'org' + str(year)
                }

    with db.session_scope() as session:

        org_data = list(_generate_org_data())
        sub_org_data = list(_generate_sub_org_data())

        nren_dict = {nren_name: model.NREN(name=nren_name)
                     for nren_name in set(d['nren'] for d in [*org_data, *sub_org_data])}
        session.add_all(nren_dict.values())

        for org in org_data:
            nren = nren_dict[org["nren"]]
            year = org["year"]
            name = org["name"]

            session.add(
                model.ParentOrganization(
                    nren=nren,
                    year=year,
                    organization=name
                )
            )
        for sub_org in sub_org_data:
            nren = nren_dict[sub_org["nren"]]
            year = sub_org["year"]
            name = sub_org["name"]
            role = sub_org["role"]

            session.add(
                model.SubOrganization(
                    nren=nren,
                    year=year,
                    organization=name,
                    role=role
                )
            )
        session.commit()


@pytest.fixture
def test_ec_project_data():
    def _generate_ec_project_data():
        for nren in ["nren" + str(i) for i in range(1, 50)]:
            for year in range(2016, 2021):
                yield {
                    'nren': nren,
                    'year': year,
                    'project': 'ec_project1',
                }
                if int(nren[-1]) % 2 == 0:
                    yield {
                        'nren': nren,
                        'year': year,
                        'project': 'ec_project2',
                    }

    with db.session_scope() as session:

        ec_project_data = list(_generate_ec_project_data())

        nren_dict = {nren_name: model.NREN(name=nren_name)
                     for nren_name in set(d['nren'] for d in ec_project_data)}
        session.add_all(nren_dict.values())

        for ec_project in ec_project_data:
            nren = nren_dict[ec_project["nren"]]
            year = ec_project["year"]
            project = ec_project["project"]

            session.add(
                model.ECProject(
                    nren=nren,
                    year=year,
                    project=project
                )
            )
        session.commit()