graphPaname/tests/preprocessing_test.py

from os import remove

from pandas import DataFrame
from requests import get

from app.preprocessing import create_dataframe
from app.data_request import request_dataset
from constants import COLUMNS, DATASETS, FILES, URL


def test_dataset_request():
    """
    Checks that the datasets URLs are reachable
    """
    for dataset in DATASETS:
        response = get(URL.format(dataset))
        assert response.status_code == 200


def test_dataframe_creation():
    """
    Verifes that the DataFrames are created and filtered properly
    """
    for dataset in DATASETS:
        request_dataset(dataset)
        df = create_dataframe(dataset)
        remove(FILES[dataset])
        assert isinstance(df, DataFrame)
        assert all(df.columns == COLUMNS[dataset])