diff --git a/flight_analysis.py b/flight_analysis.py index ec48a55..469e91d 100644 --- a/flight_analysis.py +++ b/flight_analysis.py @@ -10,7 +10,6 @@ import numpy as np import pandas as pd from datetime import timedelta, datetime -import json import configparser from src.google_flight_analysis.scrape import Scrape @@ -22,21 +21,10 @@ config.read(os.path.join(os.path.dirname(__file__), "config.ini")) -def get_routes_from_config(): - """ - Returns a list of routes from the config file. - """ - routes = [] - for route in config["routes"]: - routes.append(json.loads(config["routes"][route])) - - return routes - - if __name__ == "__main__": # 1. scrape routes - routes = get_routes_from_config() + routes = utils.get_routes_from_config(config) # compute number of total scrapes n_total_scrapes = sum([x[2] for x in routes]) diff --git a/tests/test_scrape.py b/tests/test_scrape.py index 7858206..79d24e7 100644 --- a/tests/test_scrape.py +++ b/tests/test_scrape.py @@ -1,22 +1,34 @@ import pytest import pandas as pd -from selenium import webdriver import numpy as np +from datetime import datetime, timedelta -from src.google_flight_analysis.scrape import * +from src.google_flight_analysis.scrape import Scrape +from src.google_flight_analysis.database import Database -test_df = pd.read_csv("./outputs/MUC_JFK_test.csv") +import private.private as private +import configparser -def test_dataframe_type(): - assert isinstance(test_df, pd.DataFrame) - -def test_dataframe_shape(): - assert np.array_equal(test_df.shape, np.array([241, 16])) -# def test_chromedriver_found(): -# chrome_driver = webdriver.Chrome(executable_path="/usr/bin/google-chrome") -# assert chrome_driver is not None +def test_database_connection(): + db = Database(db_host=private.DB_HOST, db_name=private.DB_NAME, db_user=private.DB_USER, db_pw=private.DB_PW, db_table=private.DB_TABLE) + try: + conn = db.connect_to_postgresql() + except ConnectionError as e: + assert False, e + + +def test_dataset_generation(): + ten_days_ahead = (datetime.today() + timedelta(5)).strftime("%Y-%m-%d") + scrape_obj = Scrape("MUC", "FCO", ten_days_ahead) + scrape_obj.run_scrape() + assert isinstance(scrape_obj.data, pd.DataFrame) -# def test_dataset_generation(): -# scrape_obj = Scrape("MUC", "FCO", "2023-11-22") -# assert isinstance(scrape_obj.data, pd.DataFrame) +def test_config_file(): + try: + config = configparser.ConfigParser() + config.read("config.ini") + except Exception as e: + assert False, e + + \ No newline at end of file diff --git a/utils.py b/utils.py index 7a5daae..63995b9 100644 --- a/utils.py +++ b/utils.py @@ -1,6 +1,7 @@ import os import logging import logging.config +import json # logging LOGS_PATH = os.path.join(os.path.dirname(__file__), "logs") @@ -42,4 +43,15 @@ def setup_logger(logger_name): create_logs_folder() logging.getLogger('WDM').setLevel(logging.NOTSET) # suppress WDM (Webdrive Manager) logs logging.config.dictConfig(LOGGING_CONFIG_DICT) - return logging.getLogger(logger_name) \ No newline at end of file + return logging.getLogger(logger_name) + + +def get_routes_from_config(config_obj): + """ + Returns a list of routes from the config file. + """ + routes = [] + for route in config_obj["routes"]: + routes.append(json.loads(config_obj["routes"][route])) + + return routes \ No newline at end of file