Compare commits

15 Commits

Author SHA1 Message Date
Jake-Pullen 5af82e5753 Bug fix no more request limit (#18)
* added tests

* Removed method no longer used due to YNAB api Changes
2025-04-04 18:49:44 +01:00
Jake d155a4c907 Fixed an issue with saving accounts data 2025-02-08 13:59:21 +00:00
Jake-Pullen 2b60d6af10 Merge pull request #16 from Jake-Pullen/document_and_tidy
Document and tidy
2024-09-04 11:51:53 +01:00
Jake Pullen 727d483e62 updated ERD 2024-09-04 11:46:35 +01:00
Jake Pullen c97a169637 added github link 2024-08-30 08:26:18 +01:00
Jake Pullen 975f0df22b Handle missing data warehouse files and bad join in dash_app 2024-08-29 11:31:35 +01:00
Jake Pullen 91d67896d1 Refactor join conditions in dash_app
fix is weekday issue, making fridays a weekend
update ERD
2024-08-29 11:02:15 +01:00
Jake Pullen bd0ebd38e9 tidying up facts and dims
removing duplicated columns
2024-08-29 10:39:55 +01:00
Jake Pullen 9e7ff808a5 tidied up dimensions and removed duplicated columns 2024-08-29 09:23:30 +01:00
Jake-Pullen d999a8175c Merge pull request #8 from Jake-Pullen/feature/add_visuals
Feature/add visuals
2024-08-28 15:26:01 +01:00
Jake Pullen 845f6a28cc separation of areas 2024-08-28 12:54:01 +01:00
Jake Pullen 7b80b52998 changes to make dash app work 2024-08-27 15:12:44 +01:00
Jake Pullen 173c0594a8 Merge branch 'main' into feature/add_visuals 2024-08-27 08:47:01 +01:00
Jake Pullen 201f8eb2c9 more dash workings 2024-08-11 10:42:05 +01:00
Jake Pullen 3504641643 understanding dash 2024-08-10 21:47:08 +01:00
17 changed files with 871 additions and 234 deletions
+2
View File
@@ -8,3 +8,5 @@ __pycache__/*
*/__pycache__/* */__pycache__/*
*.pbix *.pbix
/logs/* /logs/*
.vscode/*
*.coverage
View File
+3
View File
@@ -11,3 +11,6 @@ CONFLICT = 9
MOVE_FILE_ERROR = 10 MOVE_FILE_ERROR = 10
DUPLICATE_RESOLUTION_ERROR = 11 DUPLICATE_RESOLUTION_ERROR = 11
UNIQUE_ID_NOT_FOUND = 12 UNIQUE_ID_NOT_FOUND = 12
NO_DATA_PRODUCED = 13
MISSING_DATA_FILES = 14
BAD_JOIN = 15
+161
View File
@@ -0,0 +1,161 @@
'''Module to create a Dash app that displays visualizations of YNAB data.'''
import polars as pl
import plotly.express as px
from dash import Dash, html, dcc
import dash_bootstrap_components as dbc
import pandas as pd
import logging
import sys
import config.exit_codes as ec
try:
accounts = pl.read_parquet('data/warehouse/accounts.parquet')
categories = pl.read_parquet('data/warehouse/categories.parquet')
dates = pl.read_parquet('data/warehouse/dates.parquet')
payees = pl.read_parquet('data/warehouse/payees.parquet')
scheduled_transactions = pl.read_parquet('data/warehouse/scheduled_transactions.parquet')
transactions = pl.read_parquet('data/warehouse/transactions.parquet')
except FileNotFoundError:
logging.error('Data warehouse files not found. Run the data pipeline to create them.')
sys.exit(ec.MISSING_DATA_FILES)
try:
# Join transactions with accounts, categories, and payees to create a master DataFrame
master_transactions = transactions.join(categories, left_on='category_id', right_on='category_id', suffix='_category')\
.join(accounts, left_on='account_id', right_on='account_id', suffix='_account')\
.join(payees, left_on='payee_id', right_on='payee_id', suffix='_payee')\
.join(dates, left_on='transaction_date', right_on='date_id', suffix='_date')
except Exception as e:
logging.error(f'Error joining DataFrames: {e}')
sys.exit(ec.BAD_JOIN)
# Create aggregations
spend_per_day = master_transactions.sql('''
SELECT
date,
year,
month,
day,
ABS(SUM(transaction_amount)) as total
FROM self
WHERE category_name != 'Inflow: Ready to Assign'
GROUP BY date, year, month, day
ORDER BY date DESC
'''
)
spend_per_category = master_transactions.sql('''
SELECT
category_name,
ABS(SUM(transaction_amount)) as total
FROM self
WHERE category_name != 'Inflow: Ready to Assign'
GROUP BY category_name
ORDER BY total DESC
'''
)
spend_per_payee = master_transactions.sql('''
SELECT
payee_name,
ABS(SUM(transaction_amount)) as total
FROM self
WHERE payee_name != 'Starting Balance'
AND transaction_amount < 0
GROUP BY payee_name
ORDER BY total DESC
'''
)
# Convert DataFrame to list of dictionaries
spend_per_day_data = spend_per_day.to_dicts()
spend_per_category_data = spend_per_category.to_dicts()
spend_per_payee_data = spend_per_payee.to_dicts()
# Convert list of dictionaries to Pandas DataFrame
spend_per_day_df = pd.DataFrame(spend_per_day_data)
spend_per_category_df = pd.DataFrame(spend_per_category_data)
spend_per_payee_df = pd.DataFrame(spend_per_payee_data)
spend_per_day_line = px.line(spend_per_day_df, x="date", y="total")
spend_per_day_line.update_layout(
plot_bgcolor='black',
paper_bgcolor='black',
font_color='white'
)
spend_per_category_bar = px.bar(spend_per_category_df, x="category_name", y="total")
spend_per_category_bar.update_layout(
plot_bgcolor='black',
paper_bgcolor='black',
font_color='white'
)
spend_per_payee_bar = px.bar(spend_per_payee_df, x="payee_name", y="total")
spend_per_payee_bar.update_layout(
plot_bgcolor='black',
paper_bgcolor='black',
font_color='white'
)
# Initialize the app with a dark theme
app = Dash(external_stylesheets=[dbc.themes.DARKLY])
# App layout
app.layout = dbc.Container(
[
dbc.Row(
dbc.Col(
html.Div("Data Pipeline For YNAB, Preview Visualisations",
className="text-center text-light"),
width=12
)
),
dbc.Row(
[
dbc.Col(
dbc.Card(
dbc.CardBody(
[
html.H4("Spend Per Day", className="card-title"),
dcc.Graph(figure=spend_per_day_line)
]
),
className="mb-4"
),
width=12
)
]
),
dbc.Row(
[
dbc.Col(
dbc.Card(
dbc.CardBody(
[
html.H4("Spend Per Category", className="card-title"),
dcc.Graph(figure=spend_per_category_bar)
]
),
className="mb-4"
),
width=6
),
dbc.Col(
dbc.Card(
dbc.CardBody(
[
html.H4("Spend Per Payee", className="card-title"),
dcc.Graph(figure=spend_per_payee_bar)
]
),
className="mb-4"
),
width=6
)
]
)
],
fluid=True
)
+17 -7
View File
@@ -34,23 +34,29 @@ erDiagram
} }
DATES { DATES {
int date_id string date_id
string date date date
int year int year
int month int month
int day int day
boolean is_weekday
int weekday
} }
TRANSACTIONS { TRANSACTIONS {
int transaction_id str transaction_id
int account_id int account_id
int category_id int category_id
int payee_id int payee_id
int date_id int transaction_date
decimal amount decimal amount
boolean cleared boolean cleared
boolean approved boolean approved
boolean deleted boolean deleted
string memo
string flag_color
str transfer_account_id
} }
SCHEDULED_TRANSACTIONS { SCHEDULED_TRANSACTIONS {
@@ -58,10 +64,14 @@ erDiagram
int account_id int account_id
int category_id int category_id
int payee_id int payee_id
int date_id str date_first
str date_next
decimal amount decimal amount
string frequency string frequency
boolean deleted boolean deleted
text memo
string flag_color
str transfer_account_id
} }
TRANSACTIONS ||--o{ ACCOUNTS : "belongs to" TRANSACTIONS ||--o{ ACCOUNTS : "belongs to"
@@ -71,6 +81,6 @@ erDiagram
SCHEDULED_TRANSACTIONS ||--o{ ACCOUNTS : "belongs to" SCHEDULED_TRANSACTIONS ||--o{ ACCOUNTS : "belongs to"
SCHEDULED_TRANSACTIONS ||--o{ CATEGORIES : "belongs to" SCHEDULED_TRANSACTIONS ||--o{ CATEGORIES : "belongs to"
SCHEDULED_TRANSACTIONS ||--o{ PAYEES : "belongs to" SCHEDULED_TRANSACTIONS ||--o{ PAYEES : "belongs to"
SCHEDULED_TRANSACTIONS ||--o{ DATES : "scheduled on" SCHEDULED_TRANSACTIONS ||--o{ DATES : "First Scheduled"
SCHEDULED_TRANSACTIONS ||--o{ DATES : "Next Scheduled"
``` ```
+1 -1
View File
@@ -25,7 +25,7 @@ For the `BUDGET_ID`, you can get it from the URL of your budget page on the YNAB
### Clone the repository ### Clone the repository
```bash ```bash
git clone #link tbc git clone https://github.com/Jake-Pullen/data_pipeline_for_YNAB.git
``` ```
### Install dependencies ### Install dependencies
+4
View File
@@ -28,3 +28,7 @@ The Data Warehouse is the data after it has been aggregated and transformed. It
## Processed Archive ## Processed Archive
The Processed Archive is the data after it has been processed and stored in the base tables. It is the raw json files in the `data/processed/` directory with a folder for each entity and file for each load that has been processed. The Processed Archive is the data after it has been processed and stored in the base tables. It is the raw json files in the `data/processed/` directory with a folder for each entity and file for each load that has been processed.
## Visualisation datasets
When preparing the data for visualisation, we create dataframes in memory that are used to create the visualisations. These are not stored on disk.
+29 -36
View File
@@ -8,10 +8,7 @@ import logging.config
import logging.handlers import logging.handlers
import config.exit_codes as ec import config.exit_codes as ec
from pipeline.ingest import Ingest from pipeline.pipeline_main import pipeline_main
from pipeline.raw_to_base import RawToBase
from pipeline.dimensions import DimAccounts, DimCategories, DimPayees, DimDate
from pipeline.facts import FactTransactions, FactScheduledTransactions
def set_up_logging(): def set_up_logging():
try: try:
@@ -27,6 +24,18 @@ def set_up_logging():
queue_handler.listener.start() queue_handler.listener.start()
atexit.register(queue_handler.listener.stop) atexit.register(queue_handler.listener.stop)
def load_config():
try:
with open('config/config.yaml', 'r') as file:
config = yaml.safe_load(file)
return config
except FileNotFoundError:
logging.error('config.yaml file not found')
sys.exit(ec.MISSING_CONFIG_FILE)
except yaml.YAMLError as e:
logging.error(f'Error loading config.yaml: {e}')
sys.exit(ec.CORRUPTED_CONFIG_FILE)
logger = logging.getLogger("data_pipeline_for_ynab") logger = logging.getLogger("data_pipeline_for_ynab")
os.makedirs('logs', exist_ok=True) os.makedirs('logs', exist_ok=True)
set_up_logging() set_up_logging()
@@ -37,41 +46,25 @@ dotenv.load_dotenv()
API_TOKEN = os.getenv('API_TOKEN') API_TOKEN = os.getenv('API_TOKEN')
BUDGET_ID = os.getenv('BUDGET_ID') BUDGET_ID = os.getenv('BUDGET_ID')
def main(): if not API_TOKEN or not BUDGET_ID:
if not API_TOKEN or not BUDGET_ID: logging.error('API_TOKEN or BUDGET_ID is not set in .env file')
logging.error('API_TOKEN or BUDGET_ID is not set in .env file') sys.exit(ec.MISSING_ENV_VARS)
sys.exit(ec.MISSING_ENV_VARS)
try:
with open('config/config.yaml', 'r') as file:
config = yaml.safe_load(file)
except FileNotFoundError:
logging.error('config.yaml file not found')
sys.exit(ec.MISSING_CONFIG_FILE)
except yaml.YAMLError as e:
logging.error(f'Error loading config.yaml: {e}')
sys.exit(ec.CORRUPTED_CONFIG_FILE)
config['API_TOKEN'] = API_TOKEN
config['BUDGET_ID'] = BUDGET_ID
logging.info('Starting data pipeline')
Ingest(config)
RawToBase(config)
DimAccounts(config)
DimCategories(config)
DimPayees(config)
DimDate(config)
FactTransactions(config)
FactScheduledTransactions(config)
logging.info('Data pipeline completed successfully')
sys.exit(ec.SUCCESS)
if __name__ == '__main__': if __name__ == '__main__':
config = load_config()
config['API_TOKEN'] = API_TOKEN
config['BUDGET_ID'] = BUDGET_ID
try: try:
main() pipeline_main(config)
# Check if the data was successfully created
data_exists = os.path.exists('data/processed') and os.listdir('data/processed')
if data_exists:
from dash_app import app
app.run() # debug=True
else:
logging.error('Data pipeline did not produce any data. Dash app will not run.')
sys.exit(ec.NO_DATA_PRODUCED)
except SystemExit as e: except SystemExit as e:
exit_code = e.code exit_code = e.code
if exit_code == ec.SUCCESS: if exit_code == ec.SUCCESS:
+78 -52
View File
@@ -22,47 +22,55 @@ class DimAccounts(Dimensions):
def transform(self): def transform(self):
# Read the parquet file into a polars DataFrame # Read the parquet file into a polars DataFrame
try: try:
accounts_df = pl.read_parquet(self.file_path) source_accounts = pl.read_parquet(self.file_path)
except Exception as e: except Exception as e:
logging.error(f"Failed to read the base accounts parquet file: {e}") logging.error(f"Failed to read the base accounts parquet file: {e}")
return return
# Transform the DataFrame
logging.info("Transforming the accounts DataFrame") logging.info("Transforming the accounts DataFrame")
try: try:
accounts_df = ( base_accounts = (
accounts_df source_accounts.select([
.with_columns([ "id",
pl.col("id").alias("account_id"), "name",
pl.col("name").alias("account_name"), "type",
pl.col("type").alias("account_type"), "on_budget",
pl.col("on_budget").alias("on_budget"), "closed",
pl.col("closed").alias("closed"), "note",
pl.col("note").alias("note"), "balance",
pl.col("balance").alias("balance"), "cleared_balance",
pl.col("cleared_balance").alias("cleared_balance"), "uncleared_balance",
pl.col("uncleared_balance").alias("uncleared_balance"), "deleted"
pl.col("deleted").alias("deleted"),
])
.with_columns([
pl.col("note").fill_null("unknown"),
(pl.col("balance") / 100).alias("balance"),
(pl.col("cleared_balance") / 100).alias("cleared_balance"),
(pl.col("uncleared_balance") / 100).alias("uncleared_balance"),
])
.drop([
"transfer_payee_id", "direct_import_linked", "direct_import_in_error",
"last_reconciled_at", "debt_original_balance", "debt_interest_rates",
"debt_minimum_payments", "debt_escrow_amounts", "ingestion_date"
]) ])
) )
except Exception as e:
logging.error(f"Failed to select columns from the categories DataFrame: {e}")
return
try:
add_accounts_prefix = base_accounts.with_columns([
pl.col("id").alias("account_id"),
pl.col("name").alias("account_name"),
pl.col("type").alias("account_type")
])
fill_accounts_null_values = add_accounts_prefix.with_columns([
pl.col('note').fill_null('none')
])
fix_accounts_values = fill_accounts_null_values.with_columns([
(pl.col("balance") / 1000).alias("balance"),
(pl.col("cleared_balance") / 1000).alias("cleared_balance"),
(pl.col("uncleared_balance") / 1000).alias("uncleared_balance"),
])
drop_accounts_columns = fix_accounts_values.drop([
"id", "name", "type"
])
except Exception as e: except Exception as e:
logging.error(f"Failed to transform the accounts DataFrame: {e}") logging.error(f"Failed to transform the accounts DataFrame: {e}")
return return
# Write the DataFrame to a new parquet file
logging.info("Writing the transformed accounts DataFrame to parquet file") logging.info("Writing the transformed accounts DataFrame to parquet file")
try: try:
accounts_df.write_parquet(self.config['warehouse_data_path'] + '/accounts.parquet') drop_accounts_columns.write_parquet(self.config['warehouse_data_path'] + '/accounts.parquet')
except Exception as e: except Exception as e:
logging.error(f"Failed to write the transformed accounts DataFrame to parquet file: {e}") logging.error(f"Failed to write the transformed accounts DataFrame to parquet file: {e}")
return return
@@ -74,15 +82,14 @@ class DimCategories(Dimensions):
self.transform() self.transform()
def transform(self): def transform(self):
# Read the parquet file into a polars DataFrame
try: try:
categories_df = pl.read_parquet(self.file_path) source_categories = pl.read_parquet(self.file_path)
except Exception as e: except Exception as e:
logging.error(f"Failed to read the base categories parquet file: {e}") logging.error(f"Failed to read the base categories parquet file: {e}")
return return
logging.info("Transforming the categories DataFrame") logging.info("Transforming the categories DataFrame")
try: try:
categories_df = categories_df.select([ base_categories = source_categories.select([
'id', 'id',
'name', 'name',
'category_group_name', 'category_group_name',
@@ -98,25 +105,28 @@ class DimCategories(Dimensions):
return return
try: try:
# Rename the columns add_categories_prefix = base_categories.with_columns([
categories_df = categories_df.with_columns(pl.col('id').alias('category_id')) pl.col('id').alias('category_id'),
categories_df = categories_df.with_columns(pl.col('name').alias('category_name')) pl.col('name').alias('category_name')
])
# Fill null values in the note column fill_null_category_values = add_categories_prefix.with_columns([
categories_df = categories_df.with_columns(pl.col('note').fill_null('unknown')) pl.col('note').fill_null('none')
])
# Convert the balance, budgeted, and activity columns to decimal fix_categories_values = fill_null_category_values.with_columns([
categories_df = categories_df.with_columns(pl.col('balance') / 100) (pl.col('balance') / 1000),
categories_df = categories_df.with_columns(pl.col('budgeted') / 100) (pl.col('budgeted') / 1000),
categories_df = categories_df.with_columns(pl.col('activity') / 100) (pl.col('activity') / 1000)
])
drop_categories_columns = fix_categories_values.drop([
'id', 'name'
])
except Exception as e: except Exception as e:
logging.error(f"Failed to transform the categories DataFrame: {e}") logging.error(f"Failed to transform the categories DataFrame: {e}")
return return
# Write the DataFrame to a new parquet file
logging.info("Writing the transformed categories DataFrame to parquet file") logging.info("Writing the transformed categories DataFrame to parquet file")
try: try:
categories_df.write_parquet(self.config['warehouse_data_path'] + '/categories.parquet') drop_categories_columns.write_parquet(self.config['warehouse_data_path'] + '/categories.parquet')
except Exception as e: except Exception as e:
logging.error(f"Failed to write the transformed categories DataFrame to parquet file: {e}") logging.error(f"Failed to write the transformed categories DataFrame to parquet file: {e}")
return return
@@ -128,15 +138,14 @@ class DimPayees(Dimensions):
self.transform() self.transform()
def transform(self): def transform(self):
# Read the parquet file into a polars DataFrame
try: try:
payees_df = pl.read_parquet(self.file_path) source_payees = pl.read_parquet(self.file_path)
except Exception as e: except Exception as e:
logging.error(f"Failed to read the base payees parquet file: {e}") logging.error(f"Failed to read the base payees parquet file: {e}")
return return
logging.info("Transforming the payees DataFrame") logging.info("Transforming the payees DataFrame")
try: try:
payees_df = payees_df.select([ base_payees = source_payees.select([
'id', 'id',
'name', 'name',
'deleted' 'deleted'
@@ -144,10 +153,15 @@ class DimPayees(Dimensions):
except Exception as e: except Exception as e:
logging.error(f"Failed to select columns from the payees DataFrame: {e}") logging.error(f"Failed to select columns from the payees DataFrame: {e}")
return return
try: try:
# Rename the columns add_payees_prefix = base_payees.with_columns([
payees_df = payees_df.with_columns(pl.col('id').alias('payee_id')) pl.col('id').alias('payee_id'),
payees_df = payees_df.with_columns(pl.col('name').alias('payee_name')) pl.col('name').alias('payee_name')
])
drop_payees_columns = add_payees_prefix.drop([
'id', 'name'
])
except Exception as e: except Exception as e:
logging.error(f"Failed to rename columns in the payees DataFrame: {e}") logging.error(f"Failed to rename columns in the payees DataFrame: {e}")
return return
@@ -155,7 +169,7 @@ class DimPayees(Dimensions):
# Write the DataFrame to a new parquet file # Write the DataFrame to a new parquet file
logging.info("Writing the transformed payees DataFrame to parquet file") logging.info("Writing the transformed payees DataFrame to parquet file")
try: try:
payees_df.write_parquet(self.config['warehouse_data_path'] + '/payees.parquet') drop_payees_columns.write_parquet(self.config['warehouse_data_path'] + '/payees.parquet')
except Exception as e: except Exception as e:
logging.error(f"Failed to write the transformed payees DataFrame to parquet file: {e}") logging.error(f"Failed to write the transformed payees DataFrame to parquet file: {e}")
return return
@@ -186,11 +200,23 @@ class DimDate(Dimensions):
try: try:
# Create a new column to indicate if the date is a weekday or weekend # Create a new column to indicate if the date is a weekday or weekend
dates_df = dates_df.with_columns([ dates_df = dates_df.with_columns([
(pl.col('weekday') < 5).alias('is_weekday') # True for weekdays (Monday to Friday), False for weekends (Saturday and Sunday) (pl.col('weekday') < 6).alias('is_weekday') # True for weekdays (Monday to Friday), False for weekends (Saturday and Sunday)
]) ])
except Exception as e: except Exception as e:
logging.error(f"Failed to create a new column to indicate if the date is a weekday or weekend: {e}") logging.error(f"Failed to create a new column to indicate if the date is a weekday or weekend: {e}")
return return
# Create a primary key by concatenating year, month, and day with no separators
try:
dates_df = dates_df.with_columns([
(pl.col('year').cast(pl.Utf8) +
pl.col('month').cast(pl.Utf8).str.zfill(2) +
pl.col('day').cast(pl.Utf8).str.zfill(2)
).alias('date_id')
])
except Exception as e:
logging.error(f"Failed to create the primary key column: {e}")
return
# Write the DataFrame to a new parquet file # Write the DataFrame to a new parquet file
logging.info("Writing the transformed dates DataFrame to parquet file") logging.info("Writing the transformed dates DataFrame to parquet file")
try: try:
+93 -59
View File
@@ -18,49 +18,68 @@ class FactTransactions(Facts):
self.transform() self.transform()
def transform(self): def transform(self):
# Read the parquet file into a polars DataFrame
try: try:
transactions_df = pl.read_parquet(self.file_path) source_transactions = pl.read_parquet(self.file_path)
except FileNotFoundError: except FileNotFoundError:
logging.error("The transactions DataFrame does not exist") logging.error("The transactions DataFrame does not exist")
return return
# Transform the DataFrame try:
base_transactions = source_transactions.select([
"id",
"date",
"amount",
"memo",
"cleared",
"approved",
"flag_color",
"account_id",
"payee_id",
"category_id",
"transfer_account_id"
])
except Exception as e:
logging.error(f"Failed to select columns from the transactions DataFrame: {e}")
return
logging.info("Transforming the transactions DataFrame") logging.info("Transforming the transactions DataFrame")
try: try:
transactions_df = ( resolve_transaction_dates = base_transactions.with_columns([
transactions_df pl.col("date").str.strptime(pl.Date, format="%Y-%m-%d").alias("date")
.with_columns([ ])
pl.col("id").alias("transaction_id"), except Exception as e:
pl.col("date").alias("transaction_date"), logging.error(f"Failed to covert the date to date format: {e}")
pl.col("amount").alias("transaction_amount"), return
pl.col("memo").alias("transaction_memo"),
pl.col("cleared").alias("transaction_cleared"), try:
pl.col("approved").alias("transaction_approved"), add_transaction_prefix = resolve_transaction_dates.with_columns([
pl.col("flag_color").alias("transaction_flag_color"), pl.col("id").alias("transaction_id"),
pl.col("account_id").alias("account_id"), (pl.col("date").dt.year().cast(pl.Utf8) +
pl.col("payee_id").alias("payee_id"), pl.col("date").dt.month().cast(pl.Utf8).str.zfill(2) +
pl.col("category_id").alias("category_id"), pl.col("date").dt.day().cast(pl.Utf8).str.zfill(2)).alias("transaction_date"),
pl.col("transfer_account_id").alias("transfer_account_id"), ])
]) fix_transaction_nulls = add_transaction_prefix.with_columns([
.with_columns([ pl.col("memo").fill_null("none"),
pl.col("memo").fill_null("unknown"), pl.col("flag_color").fill_null("none"),
(pl.col("amount") / 100).alias("transaction_amount"), pl.col("transfer_account_id").fill_null("none"),
]) pl.col("category_id").fill_null("none"),
.drop([ ])
"transfer_transaction_id", "matched_transaction_id", "import_id", fix_transaction_values = fix_transaction_nulls.with_columns([
"subtransactions", "deleted","flag_name","account_name", (pl.col("amount") / 1000).alias("transaction_amount")
"payee_name","category_name","import_payee_name","import_payee_name_original", ])
"debt_transaction_type","ingestion_date" drop_transaction_columns = fix_transaction_values.drop([
]) "id", "date", "amount"
) ])
except Exception as e: except Exception as e:
logging.error(f"Failed to transform the transactions DataFrame: {e}") logging.error(f"Failed to transform the transactions DataFrame: {e}")
return return
# Write the DataFrame to a new parquet file # Write the DataFrame to a new parquet file
logging.info("Writing the transformed transactions DataFrame to parquet file") logging.info("Writing the transformed transactions DataFrame to parquet file")
try: try:
transactions_df.write_parquet(self.config['warehouse_data_path'] + '/transactions.parquet') drop_transaction_columns.write_parquet(
self.config['warehouse_data_path'] + '/transactions.parquet'
)
except Exception as e: except Exception as e:
logging.error(f"Failed to write the transformed transactions DataFrame: {e}") logging.error(f"Failed to write the transformed transactions DataFrame: {e}")
@@ -71,46 +90,61 @@ class FactScheduledTransactions(Facts):
self.transform() self.transform()
def transform(self): def transform(self):
# Read the parquet file into a polars DataFrame
try: try:
scheduled_transactions_df = pl.read_parquet(self.file_path) source_scheduled = pl.read_parquet(self.file_path)
except FileNotFoundError: except FileNotFoundError:
logging.error("The scheduled transactions DataFrame does not exist") logging.error("The scheduled transactions DataFrame does not exist")
return return
# Transform the DataFrame try:
base_scheduled = source_scheduled.select([
"id",
"date_first",
"date_next",
"frequency",
"amount",
"memo",
"flag_color",
"account_id",
"payee_id",
"category_id",
"transfer_account_id"
])
except Exception as e:
logging.error(f"Failed to select columns from the scheduled transactions DataFrame: {e}")
return
try:
resolve_scheduled_dates = base_scheduled.with_columns([
pl.col("date_first").str.strptime(pl.Date, format="%Y-%m-%d").alias("date_first"),
pl.col("date_next").str.strptime(pl.Date, format="%Y-%m-%d").alias("date_next")
])
except Exception as e:
logging.error(f"Failed to covert the date to date format: {e}")
return
logging.info("Transforming the scheduled transactions DataFrame") logging.info("Transforming the scheduled transactions DataFrame")
try: try:
scheduled_transactions_df = ( add_scheduled_prefix = resolve_scheduled_dates.with_columns([
scheduled_transactions_df pl.col("id").alias("scheduled_transaction_id")
.with_columns([ ])
pl.col("id").alias("scheduled_transaction_id"), fix_sheduled_nulls = add_scheduled_prefix.with_columns([
pl.col("date_first").alias("scheduled_transaction_first_date"), pl.col("memo").fill_null("none"),
pl.col("date_next").alias("scheduled_transaction_next_date"), pl.col("flag_color").fill_null("none"),
pl.col("frequency").alias("scheduled_transaction_frequency"), pl.col("transfer_account_id").fill_null("none"),
pl.col("amount").alias("scheduled_transaction_amount"), pl.col("category_id").fill_null("none"),
pl.col("memo").alias("scheduled_transaction_memo"), ])
pl.col("flag_color").alias("scheduled_transaction_flag_color"), fix_scheduled_values = fix_sheduled_nulls.with_columns([
pl.col("account_id").alias("account_id"), (pl.col("amount") / 1000).alias("scheduled_transaction_amount"),
pl.col("payee_id").alias("payee_id"), ])
pl.col("category_id").alias("category_id"), drop_scheduled_columns = fix_scheduled_values.drop([
pl.col("transfer_account_id").alias("transfer_account_id"), "id", "amount"
]) ])
.with_columns([
pl.col("memo").fill_null("unknown"),
(pl.col("amount") / 100).alias("scheduled_transaction_amount"),
])
.drop([
"subtransactions", "deleted","flag_name","account_name",
"payee_name","category_name","ingestion_date"
])
)
except Exception as e: except Exception as e:
logging.error(f"Failed to transform the scheduled transactions DataFrame: {e}") logging.error(f"Failed to transform the scheduled transactions DataFrame: {e}")
return return
# Write the DataFrame to a new parquet file
logging.info("Writing the transformed scheduled transactions DataFrame to parquet file") logging.info("Writing the transformed scheduled transactions DataFrame to parquet file")
try: try:
scheduled_transactions_df.write_parquet(self.config['warehouse_data_path'] + '/scheduled_transactions.parquet') drop_scheduled_columns.write_parquet(self.config['warehouse_data_path'] + '/scheduled_transactions.parquet')
except Exception as e: except Exception as e:
logging.error(f"Failed to write the transformed scheduled transactions DataFrame: {e}") logging.error(f"Failed to write the transformed scheduled transactions DataFrame: {e}")
+26 -38
View File
@@ -4,13 +4,11 @@ import json
import logging import logging
import requests import requests
import sys import sys
import yaml
from typing import Dict, Any from typing import Dict, Any
import config.exit_codes as ec import config.exit_codes as ec
class Ingest: class Ingest:
def __init__(self, config: Dict[str, Any]): def __init__(self, config: Dict[str, Any]):
""" """
Initialize the Ingest class with the provided configuration. Initialize the Ingest class with the provided configuration.
@@ -22,19 +20,9 @@ class Ingest:
self.entities = config['entities'] self.entities = config['entities']
self.raw_data_path = config['raw_data_path'] self.raw_data_path = config['raw_data_path']
self.headers = {'Authorization': f'Bearer {self.api_token}'} self.headers = {'Authorization': f'Bearer {self.api_token}'}
self.knowledge_cache = self.load_knowledge_cache()
self.MAX_RETRIES = config['REQUESTS_MAX_RETRIES'] self.MAX_RETRIES = config['REQUESTS_MAX_RETRIES']
self.RETRY_DELAY = config['REQUESTS_RETRY_DELAY'] self.RETRY_DELAY = config['REQUESTS_RETRY_DELAY']
self.fetch_and_cache_entity_data()
def load_knowledge_cache(self) -> Dict[str, Any]:
"""
Load the knowledge cache from the file if it exists.
"""
if os.path.exists(self.knowledge_file):
with open(self.knowledge_file, 'r') as f:
return json.load(f)
return {}
def save_entity_data_to_raw(self, entity: str, data: Dict[str, Any]): def save_entity_data_to_raw(self, entity: str, data: Dict[str, Any]):
""" """
@@ -50,8 +38,18 @@ class Ingest:
with open(entity_file, 'w') as f: with open(entity_file, 'w') as f:
json.dump(data, f, indent=4) json.dump(data, f, indent=4)
except Exception as e: except Exception as e:
logging.error(f"Error saving {entity} data: {e}") logging.error(f"Failed to save data for {entity} to {entity_file}")
raise e
def load_knowledge_cache(self) -> Dict[str, Any]:
"""
Load the knowledge cache from the file if it exists.
"""
if not os.path.exists(self.knowledge_file):
os.makedirs(os.path.dirname(self.knowledge_file),exist_ok=True)
return {}
with open(self.knowledge_file, 'r') as f:
return json.load(f)
def update_server_knowledge_cache(self, entity: str, server_knowledge: Any): def update_server_knowledge_cache(self, entity: str, server_knowledge: Any):
""" """
@@ -70,24 +68,14 @@ class Ingest:
with open(self.knowledge_file, 'w') as f: with open(self.knowledge_file, 'w') as f:
json.dump(knowledge_cache, f, indent=4) json.dump(knowledge_cache, f, indent=4)
def check_rate_limit(self, response: requests.Response): knowledge_cache = self.load_knowledge_cache()
""" knowledge_cache[entity] = server_knowledge
Check and handle the rate limit based on the response headers. try:
""" with open(self.knowledge_file, 'w') as f:
rate_limit_header = response.headers.get('X-Rate-Limit') json.dump(knowledge_cache, f, indent=4)
if rate_limit_header: except Exception as e:
requests_made, limit = map(int, rate_limit_header.split('/')) logging.error(f"Failed to update knowledge cache for {entity} in {self.knowledge_file}")
remaining_requests = limit - requests_made raise e
logging.info(f"Rate Limit: {remaining_requests}/{limit} requests remaining.")
if remaining_requests < 20:
logging.warning("Approaching rate limit. Consider pausing further requests.")
# Implement pause or delay logic here if necessary
if remaining_requests == 1:
logging.error("Rate limit exceeded. ending requests here and moving on with what we have.")
return True #returning True here to break out of any more ingestions
else:
logging.warning("X-Rate-Limit header is missing.")
def handle_response(self, response) -> bool: def handle_response(self, response) -> bool:
if response.status_code == 400: if response.status_code == 400:
@@ -100,7 +88,7 @@ class Ingest:
logging.error("Forbidden. Access is denied.") logging.error("Forbidden. Access is denied.")
sys.exit(ec.FORBIDDEN) sys.exit(ec.FORBIDDEN)
elif response.status_code == 404: elif response.status_code == 404:
logging.error("Not found. The specified URI does not exist.") logging.error("Not found. The specified URL does not exist.")
sys.exit(ec.NOT_FOUND) sys.exit(ec.NOT_FOUND)
elif response.status_code == 409: elif response.status_code == 409:
logging.error("Conflict. The resource cannot be saved due to a conflict.") logging.error("Conflict. The resource cannot be saved due to a conflict.")
@@ -118,7 +106,7 @@ class Ingest:
response.raise_for_status() response.raise_for_status()
return False return False
def fetch_and_cache_entity_data(self): def start_ingestion(self):
""" """
Fetch and cache data for all entities. Fetch and cache data for all entities.
""" """
@@ -128,11 +116,13 @@ class Ingest:
logging.warning(f"Raw data exists for {entity} processing any raw data we already have.") logging.warning(f"Raw data exists for {entity} processing any raw data we already have.")
break # break here instead of continue as we dont want to update our server knowledge cache and potentially miss data. break # break here instead of continue as we dont want to update our server knowledge cache and potentially miss data.
last_knowledge = self.knowledge_cache.get(entity, 0) knowledge_cache = self.load_knowledge_cache()
last_knowledge = knowledge_cache.get(entity, 0)
#logging.debug(f'Last Knowledge of {entity}: {last_knowledge}') #logging.debug(f'Last Knowledge of {entity}: {last_knowledge}')
logging.info(f'Fetching {entity} data since last knowledge: {last_knowledge}') logging.info(f'Fetching {entity} data since last knowledge: {last_knowledge}')
url = f'{self.base_url}/{self.budget_id}/{entity}?last_knowledge_of_server={last_knowledge}' url = f'{self.base_url}/{self.budget_id}/{entity}?last_knowledge_of_server={last_knowledge}'
response = None
for attempt in range(self.MAX_RETRIES): for attempt in range(self.MAX_RETRIES):
try: try:
response = requests.get(url, headers=self.headers) response = requests.get(url, headers=self.headers)
@@ -148,6 +138,7 @@ class Ingest:
sys.exit(ec.REQUESTS_ERROR) sys.exit(ec.REQUESTS_ERROR)
data = response.json() data = response.json()
logging.debug(f'response data: {data}')
server_knowledge = data['data'].get('server_knowledge') server_knowledge = data['data'].get('server_knowledge')
logging.debug(f'{entity} new server knowledge: {server_knowledge}') logging.debug(f'{entity} new server knowledge: {server_knowledge}')
@@ -158,6 +149,3 @@ class Ingest:
self.save_entity_data_to_raw(entity, entity_data) self.save_entity_data_to_raw(entity, entity_data)
else: else:
logging.info(f"No new data for {entity}. Skipping cache update.") logging.info(f"No new data for {entity}. Skipping cache update.")
if self.check_rate_limit(response):
break # break out here and continue processing the data we have.
+25
View File
@@ -0,0 +1,25 @@
'''Module to run the data pipeline'''
import logging
from pipeline.ingest import Ingest
from pipeline.raw_to_base import RawToBase
from pipeline.dimensions import DimAccounts, DimCategories, DimPayees, DimDate
from pipeline.facts import FactTransactions, FactScheduledTransactions
def pipeline_main(config):
'''Run the data pipeline'''
logging.info('Starting data pipeline')
ingest = Ingest(config)
ingest.start_ingestion()
RawToBase(config)
DimAccounts(config)
DimCategories(config)
DimPayees(config)
DimDate(config)
FactTransactions(config)
FactScheduledTransactions(config)
logging.info('Data pipeline completed successfully')
+7 -1
View File
@@ -68,6 +68,12 @@ Then move the files back in one at a time oldest to newest and run again for eac
return False return False
modified_data = self._add_ingestion_date(entity, data, file_name) modified_data = self._add_ingestion_date(entity, data, file_name)
for index, record in enumerate(modified_data):
logging.debug(f"processing record: {record}")
filtered_record = {k: v for k, v in record.items() if not k.startswith('debt_')}
modified_data[index] = filtered_record
logging.debug(f"filtered record: {filtered_record}")
logging.debug(f"modified data: {modified_data}")
self.data[entity].append(modified_data) self.data[entity].append(modified_data)
logging.debug(f"Successfully loaded data from file: {file_path}") logging.debug(f"Successfully loaded data from file: {file_path}")
@@ -130,7 +136,7 @@ Then move the files back in one at a time oldest to newest and run again for eac
df = df.with_columns( df = df.with_columns(
pl.when(pl.col(col).is_null()) pl.when(pl.col(col).is_null())
.then(pl.lit("null")) .then(pl.lit("null"))
.otherwise(pl.col(col).map_elements(lambda x: str(x) if x is not None else "null")) .otherwise(pl.col(col).map_elements(lambda x: str(x) if x is not None else "null", return_dtype=pl.Utf8))
.alias(col) .alias(col)
) )
return df return df
+7
View File
@@ -2,3 +2,10 @@ python-dotenv
polars polars
requests requests
pyyaml pyyaml
#visualisation requirements below
dash
pandas
pyarrow
dash-bootstrap-components
# testing requirements below
pytest
+307
View File
@@ -0,0 +1,307 @@
import pytest
from unittest.mock import patch, mock_open, MagicMock
import json
import os
from typing import Dict, Any
import logging
from pipeline.ingest import Ingest
import config.exit_codes as ec
# Mock configuration for initializing the Ingest class
mock_config = {
'API_TOKEN': 'test_token',
'BUDGET_ID': 'test_budget_id',
'base_url': 'http://test_base_url',
'knowledge_file': 'data/test_knowledge_file.json',
'entities': ['entity1', 'entity2'],
'raw_data_path': 'test_raw_data_path',
'REQUESTS_MAX_RETRIES': 3,
'REQUESTS_RETRY_DELAY': 1
}
# Test for load_knowledge_cache method
def test_load_knowledge_cache_file_exists():
mock_data = {"key": "value"}
with patch('os.path.exists', return_value=True), \
patch('builtins.open', mock_open(read_data=json.dumps(mock_data))) as mock_file:
ingest_instance = Ingest(mock_config)
result = ingest_instance.load_knowledge_cache()
mock_file.assert_called_once_with(mock_config['knowledge_file'], 'r')
assert result == mock_data
def test_load_knowledge_cache_file_not_exists():
with patch('os.path.exists', return_value=False):
ingest_instance = Ingest(mock_config)
result = ingest_instance.load_knowledge_cache()
assert result == {}
# Test for save_entity_data_to_raw method
def test_save_entity_data_to_raw_success():
entity = 'entity1'
data = {"key": "value"}
current_time = '20230101123000'
directory = os.path.join(mock_config['raw_data_path'], entity)
entity_file = f'{directory}/{current_time}.json'
with patch('os.path.exists', return_value=False), \
patch('os.makedirs') as mock_makedirs, \
patch('builtins.open', mock_open()) as mock_file, \
patch('time.strftime', return_value=current_time), \
patch('logging.info') as mock_logging_info:
ingest_instance = Ingest(mock_config)
ingest_instance.save_entity_data_to_raw(entity, data)
mock_makedirs.assert_called_once_with(directory)
mock_file.assert_called_once_with(entity_file, 'w')
# Get the file handle and check the written content
handle = mock_file()
handle.write.assert_called()
written_content = ''.join(call.args[0] for call in handle.write.call_args_list)
assert written_content == json.dumps(data, indent=4)
mock_logging_info.assert_called_once_with(f"Saving {entity} data to {entity_file}")
def test_save_entity_data_to_raw_existing_directory():
entity = 'entity1'
data = {"key": "value"}
current_time = '20230101123000'
directory = os.path.join(mock_config['raw_data_path'], entity)
entity_file = f'{directory}/{current_time}.json'
with patch('os.path.exists', return_value=True), \
patch('os.makedirs') as mock_makedirs, \
patch('builtins.open', mock_open()) as mock_file, \
patch('time.strftime', return_value=current_time), \
patch('logging.info') as mock_logging_info:
ingest_instance = Ingest(mock_config)
ingest_instance.save_entity_data_to_raw(entity, data)
mock_makedirs.assert_not_called()
mock_file.assert_called_once_with(entity_file, 'w')
# Get the file handle and check the written content
handle = mock_file()
handle.write.assert_called()
written_content = ''.join(call.args[0] for call in handle.write.call_args_list)
assert written_content == json.dumps(data, indent=4)
mock_logging_info.assert_called_once_with(f"Saving {entity} data to {entity_file}")
def test_save_entity_data_to_raw_error():
entity = 'entity1'
data = {"key": "value"}
current_time = '20230101123000'
directory = os.path.join(mock_config['raw_data_path'], entity)
entity_file = f'{directory}/{current_time}.json'
with patch('os.path.exists', return_value=True), \
patch('builtins.open', mock_open()) as mock_file, \
patch('time.strftime', return_value=current_time), \
patch('logging.info') as mock_logging_info, \
patch('logging.error') as mock_logging_error:
mock_file.side_effect = Exception("Test error")
ingest_instance = Ingest(mock_config)
with pytest.raises(Exception, match="Test error"):
ingest_instance.save_entity_data_to_raw(entity, data)
mock_logging_error.assert_called_once_with(f"Failed to save data for {entity} to {entity_file}")
def test_update_server_knowledge_cache_file_exists():
entity = 'entity1'
server_knowledge = {"key": "value"}
existing_cache = {"entity2": {"key": "old_value"}}
updated_cache = {"entity2": {"key": "old_value"}, "entity1": {"key": "value"}}
with patch('builtins.open', mock_open(read_data=json.dumps(existing_cache))) as mock_file, \
patch('os.path.exists', return_value=True), \
patch('logging.error') as mock_logging_error:
ingest_instance = Ingest(mock_config)
ingest_instance.update_server_knowledge_cache(entity, server_knowledge)
mock_file.assert_called_with(mock_config['knowledge_file'], 'w')
handle = mock_file()
handle.write.assert_called()
written_content = ''.join(call.args[0] for call in handle.write.call_args_list)
assert json.loads(written_content) == updated_cache
mock_logging_error.assert_not_called()
def test_update_server_knowledge_cache_file_not_exists():
entity = 'entity1'
server_knowledge = {"key": "value"}
updated_cache = {"entity1": {"key": "value"}}
with patch('builtins.open', mock_open()) as mock_file, \
patch('os.path.exists', return_value=False), \
patch('os.makedirs') as mock_makedirs, \
patch('logging.info') as mock_logging_info, \
patch('logging.error') as mock_logging_error:
# Ensure the side_effect list has enough elements to cover all calls to open
mock_file.side_effect = [FileNotFoundError(), mock_open().return_value]
ingest_instance = Ingest(mock_config)
with pytest.raises(FileNotFoundError):
ingest_instance.update_server_knowledge_cache(entity, server_knowledge)
mock_makedirs.assert_called_once_with(os.path.dirname(mock_config['knowledge_file']), exist_ok=True)
mock_file.assert_called_with(mock_config['knowledge_file'], 'w')
mock_logging_error.assert_called_once_with(f"Failed to update knowledge cache for {entity} in {mock_config['knowledge_file']}")
def test_update_server_knowledge_cache_write_error():
entity = 'entity1'
server_knowledge = {"key": "value"}
with patch('builtins.open', mock_open()) as mock_file, \
patch('logging.error') as mock_logging_error:
mock_file.side_effect = Exception("Test error")
ingest_instance = Ingest(mock_config)
with pytest.raises(Exception, match="Test error"):
ingest_instance.update_server_knowledge_cache(entity, server_knowledge)
mock_logging_error.assert_called_once_with(f"Failed to update knowledge cache for {entity} in {mock_config['knowledge_file']}")
def test_check_rate_limit_above_threshold():
response = MagicMock()
response.headers = {'X-Rate-Limit': '10/100'}
ingest_instance = Ingest(mock_config)
result = ingest_instance.check_rate_limit(response)
assert result is None
def test_check_rate_limit_below_threshold():
response = MagicMock()
response.headers = {'X-Rate-Limit': '90/100'}
ingest_instance = Ingest(mock_config)
result = ingest_instance.check_rate_limit(response)
assert result is None
def test_check_rate_limit_exceeded():
response = MagicMock()
response.headers = {'X-Rate-Limit': '100/100'}
ingest_instance = Ingest(mock_config)
result = ingest_instance.check_rate_limit(response)
assert result is True
def test_check_rate_limit_header_missing():
response = MagicMock()
response.headers = {}
ingest_instance = Ingest(mock_config)
result = ingest_instance.check_rate_limit(response)
assert result is None
def test_handle_response_bad_request():
response = MagicMock()
response.status_code = 400
ingest_instance = Ingest(mock_config)
with pytest.raises(SystemExit) as e:
ingest_instance.handle_response(response)
assert e.type == SystemExit
assert e.value.code == ec.BAD_REQUEST
def test_handle_response_unauthorized():
response = MagicMock()
response.status_code = 401
ingest_instance = Ingest(mock_config)
with pytest.raises(SystemExit) as e:
ingest_instance.handle_response(response)
assert e.type == SystemExit
assert e.value.code == ec.UNAUTHORIZED_API_TOKEN
def test_handle_response_forbidden():
response = MagicMock()
response.status_code = 403
ingest_instance = Ingest(mock_config)
with pytest.raises(SystemExit) as e:
ingest_instance.handle_response(response)
assert e.type == SystemExit
assert e.value.code == ec.FORBIDDEN
def test_handle_response_not_found():
response = MagicMock()
response.status_code = 404
ingest_instance = Ingest(mock_config)
with pytest.raises(SystemExit) as e:
ingest_instance.handle_response(response)
assert e.type == SystemExit
assert e.value.code == ec.NOT_FOUND
def test_handle_response_conflict():
response = MagicMock()
response.status_code = 409
ingest_instance = Ingest(mock_config)
with pytest.raises(SystemExit) as e:
ingest_instance.handle_response(response)
assert e.type == SystemExit
assert e.value.code == ec.CONFLICT
def test_handle_response_too_many_requests():
response = MagicMock()
response.status_code = 429
ingest_instance = Ingest(mock_config)
result = ingest_instance.handle_response(response)
assert result is True
def test_handle_response_internal_server_error():
response = MagicMock()
response.status_code = 500
ingest_instance = Ingest(mock_config)
result = ingest_instance.handle_response(response)
assert result is True
def test_handle_response_service_unavailable():
response = MagicMock()
response.status_code = 503
ingest_instance = Ingest(mock_config)
result = ingest_instance.handle_response(response)
assert result is True
def test_handle_response_ok():
response = MagicMock()
response.status_code = 200
ingest_instance = Ingest(mock_config)
result = ingest_instance.handle_response(response)
assert result is False
if __name__ == "__main__":
pytest.main()
+71
View File
@@ -0,0 +1,71 @@
import pytest
from unittest.mock import patch, mock_open, MagicMock
import yaml
import logging
import atexit
import sys
from main import set_up_logging, load_config
import config.exit_codes as ec
# Test for set_up_logging function
def test_set_up_logging_success():
with patch('builtins.open', mock_open(read_data="handlers:\n queue_handler:\n class: logging.handlers.QueueHandler")), \
patch('yaml.safe_load', return_value={"handlers": {"queue_handler": {"class": "logging.handlers.QueueHandler"}}}), \
patch('logging.config.dictConfig') as mock_dict_config, \
patch('logging.getHandlerByName', return_value=MagicMock(listener=MagicMock(start=MagicMock(), stop=MagicMock()))), \
patch('atexit.register') as mock_atexit_register:
set_up_logging()
mock_dict_config.assert_called_once_with({"handlers": {"queue_handler": {"class": "logging.handlers.QueueHandler"}}})
mock_atexit_register.assert_called_once()
def test_set_up_logging_yaml_error():
with patch('builtins.open', mock_open(read_data="invalid_yaml")), \
patch('yaml.safe_load', side_effect=yaml.YAMLError("Error")), \
patch('logging.basicConfig') as mock_basic_config:
set_up_logging()
mock_basic_config.assert_called_once_with(level=logging.INFO)
def test_set_up_logging_no_queue_handler():
with patch('builtins.open', mock_open(read_data="handlers:\n queue_handler:\n class: logging.handlers.QueueHandler")), \
patch('yaml.safe_load', return_value={"handlers": {"queue_handler": {"class": "logging.handlers.QueueHandler"}}}), \
patch('logging.config.dictConfig') as mock_dict_config, \
patch('logging.getHandlerByName', return_value=None):
set_up_logging()
mock_dict_config.assert_called_once_with({"handlers": {"queue_handler": {"class": "logging.handlers.QueueHandler"}}})
# Test for load_config function
def test_load_config_success():
with patch('builtins.open', mock_open(read_data="key: value")), \
patch('yaml.safe_load', return_value={"key": "value"}):
config = load_config()
assert config == {"key": "value"}
def test_load_config_file_not_found():
with patch('builtins.open', side_effect=FileNotFoundError), \
patch('logging.error') as mock_logging_error, \
patch('sys.exit') as mock_sys_exit:
load_config()
mock_logging_error.assert_called_once_with('config.yaml file not found')
mock_sys_exit.assert_called_once_with(ec.MISSING_CONFIG_FILE)
def test_load_config_yaml_error():
with patch('builtins.open', mock_open(read_data="invalid_yaml")), \
patch('yaml.safe_load', side_effect=yaml.YAMLError("Error")), \
patch('logging.error') as mock_logging_error, \
patch('sys.exit') as mock_sys_exit:
load_config()
mock_logging_error.assert_called_once()
mock_sys_exit.assert_called_once_with(ec.CORRUPTED_CONFIG_FILE)