Compare commits
15 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 5af82e5753 | |||
| d155a4c907 | |||
| 2b60d6af10 | |||
| 727d483e62 | |||
| c97a169637 | |||
| 975f0df22b | |||
| 91d67896d1 | |||
| bd0ebd38e9 | |||
| 9e7ff808a5 | |||
| d999a8175c | |||
| 845f6a28cc | |||
| 7b80b52998 | |||
| 173c0594a8 | |||
| 201f8eb2c9 | |||
| 3504641643 |
@@ -8,3 +8,5 @@ __pycache__/*
|
||||
*/__pycache__/*
|
||||
*.pbix
|
||||
/logs/*
|
||||
.vscode/*
|
||||
*.coverage
|
||||
|
||||
@@ -11,3 +11,6 @@ CONFLICT = 9
|
||||
MOVE_FILE_ERROR = 10
|
||||
DUPLICATE_RESOLUTION_ERROR = 11
|
||||
UNIQUE_ID_NOT_FOUND = 12
|
||||
NO_DATA_PRODUCED = 13
|
||||
MISSING_DATA_FILES = 14
|
||||
BAD_JOIN = 15
|
||||
+161
@@ -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
@@ -34,23 +34,29 @@ erDiagram
|
||||
}
|
||||
|
||||
DATES {
|
||||
int date_id
|
||||
string date
|
||||
string date_id
|
||||
date date
|
||||
int year
|
||||
int month
|
||||
int day
|
||||
boolean is_weekday
|
||||
int weekday
|
||||
}
|
||||
|
||||
TRANSACTIONS {
|
||||
int transaction_id
|
||||
str transaction_id
|
||||
int account_id
|
||||
int category_id
|
||||
int payee_id
|
||||
int date_id
|
||||
int transaction_date
|
||||
decimal amount
|
||||
boolean cleared
|
||||
boolean approved
|
||||
boolean deleted
|
||||
string memo
|
||||
string flag_color
|
||||
str transfer_account_id
|
||||
|
||||
}
|
||||
|
||||
SCHEDULED_TRANSACTIONS {
|
||||
@@ -58,10 +64,14 @@ erDiagram
|
||||
int account_id
|
||||
int category_id
|
||||
int payee_id
|
||||
int date_id
|
||||
str date_first
|
||||
str date_next
|
||||
decimal amount
|
||||
string frequency
|
||||
boolean deleted
|
||||
text memo
|
||||
string flag_color
|
||||
str transfer_account_id
|
||||
}
|
||||
|
||||
TRANSACTIONS ||--o{ ACCOUNTS : "belongs to"
|
||||
@@ -71,6 +81,6 @@ erDiagram
|
||||
SCHEDULED_TRANSACTIONS ||--o{ ACCOUNTS : "belongs to"
|
||||
SCHEDULED_TRANSACTIONS ||--o{ CATEGORIES : "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
@@ -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
|
||||
|
||||
```bash
|
||||
git clone #link tbc
|
||||
git clone https://github.com/Jake-Pullen/data_pipeline_for_YNAB.git
|
||||
```
|
||||
|
||||
### Install dependencies
|
||||
|
||||
@@ -28,3 +28,7 @@ The Data Warehouse is the data after it has been aggregated and transformed. It
|
||||
## 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.
|
||||
|
||||
## 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.
|
||||
|
||||
@@ -8,10 +8,7 @@ import logging.config
|
||||
import logging.handlers
|
||||
|
||||
import config.exit_codes as ec
|
||||
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
|
||||
from pipeline.pipeline_main import pipeline_main
|
||||
|
||||
def set_up_logging():
|
||||
try:
|
||||
@@ -27,6 +24,18 @@ def set_up_logging():
|
||||
queue_handler.listener.start()
|
||||
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")
|
||||
os.makedirs('logs', exist_ok=True)
|
||||
set_up_logging()
|
||||
@@ -37,41 +46,25 @@ dotenv.load_dotenv()
|
||||
API_TOKEN = os.getenv('API_TOKEN')
|
||||
BUDGET_ID = os.getenv('BUDGET_ID')
|
||||
|
||||
def main():
|
||||
if not API_TOKEN or not BUDGET_ID:
|
||||
logging.error('API_TOKEN or BUDGET_ID is not set in .env file')
|
||||
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 not API_TOKEN or not BUDGET_ID:
|
||||
logging.error('API_TOKEN or BUDGET_ID is not set in .env file')
|
||||
sys.exit(ec.MISSING_ENV_VARS)
|
||||
|
||||
if __name__ == '__main__':
|
||||
config = load_config()
|
||||
config['API_TOKEN'] = API_TOKEN
|
||||
config['BUDGET_ID'] = BUDGET_ID
|
||||
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:
|
||||
exit_code = e.code
|
||||
if exit_code == ec.SUCCESS:
|
||||
|
||||
+78
-52
@@ -22,47 +22,55 @@ class DimAccounts(Dimensions):
|
||||
def transform(self):
|
||||
# Read the parquet file into a polars DataFrame
|
||||
try:
|
||||
accounts_df = pl.read_parquet(self.file_path)
|
||||
source_accounts = pl.read_parquet(self.file_path)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to read the base accounts parquet file: {e}")
|
||||
return
|
||||
|
||||
# Transform the DataFrame
|
||||
logging.info("Transforming the accounts DataFrame")
|
||||
try:
|
||||
accounts_df = (
|
||||
accounts_df
|
||||
.with_columns([
|
||||
pl.col("id").alias("account_id"),
|
||||
pl.col("name").alias("account_name"),
|
||||
pl.col("type").alias("account_type"),
|
||||
pl.col("on_budget").alias("on_budget"),
|
||||
pl.col("closed").alias("closed"),
|
||||
pl.col("note").alias("note"),
|
||||
pl.col("balance").alias("balance"),
|
||||
pl.col("cleared_balance").alias("cleared_balance"),
|
||||
pl.col("uncleared_balance").alias("uncleared_balance"),
|
||||
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"
|
||||
base_accounts = (
|
||||
source_accounts.select([
|
||||
"id",
|
||||
"name",
|
||||
"type",
|
||||
"on_budget",
|
||||
"closed",
|
||||
"note",
|
||||
"balance",
|
||||
"cleared_balance",
|
||||
"uncleared_balance",
|
||||
"deleted"
|
||||
])
|
||||
)
|
||||
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:
|
||||
logging.error(f"Failed to transform the accounts DataFrame: {e}")
|
||||
return
|
||||
# Write the DataFrame to a new parquet file
|
||||
|
||||
logging.info("Writing the transformed accounts DataFrame to parquet file")
|
||||
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:
|
||||
logging.error(f"Failed to write the transformed accounts DataFrame to parquet file: {e}")
|
||||
return
|
||||
@@ -74,15 +82,14 @@ class DimCategories(Dimensions):
|
||||
self.transform()
|
||||
|
||||
def transform(self):
|
||||
# Read the parquet file into a polars DataFrame
|
||||
try:
|
||||
categories_df = pl.read_parquet(self.file_path)
|
||||
source_categories = pl.read_parquet(self.file_path)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to read the base categories parquet file: {e}")
|
||||
return
|
||||
logging.info("Transforming the categories DataFrame")
|
||||
try:
|
||||
categories_df = categories_df.select([
|
||||
base_categories = source_categories.select([
|
||||
'id',
|
||||
'name',
|
||||
'category_group_name',
|
||||
@@ -98,25 +105,28 @@ class DimCategories(Dimensions):
|
||||
return
|
||||
|
||||
try:
|
||||
# Rename the columns
|
||||
categories_df = categories_df.with_columns(pl.col('id').alias('category_id'))
|
||||
categories_df = categories_df.with_columns(pl.col('name').alias('category_name'))
|
||||
|
||||
# Fill null values in the note column
|
||||
categories_df = categories_df.with_columns(pl.col('note').fill_null('unknown'))
|
||||
|
||||
# Convert the balance, budgeted, and activity columns to decimal
|
||||
categories_df = categories_df.with_columns(pl.col('balance') / 100)
|
||||
categories_df = categories_df.with_columns(pl.col('budgeted') / 100)
|
||||
categories_df = categories_df.with_columns(pl.col('activity') / 100)
|
||||
add_categories_prefix = base_categories.with_columns([
|
||||
pl.col('id').alias('category_id'),
|
||||
pl.col('name').alias('category_name')
|
||||
])
|
||||
fill_null_category_values = add_categories_prefix.with_columns([
|
||||
pl.col('note').fill_null('none')
|
||||
])
|
||||
fix_categories_values = fill_null_category_values.with_columns([
|
||||
(pl.col('balance') / 1000),
|
||||
(pl.col('budgeted') / 1000),
|
||||
(pl.col('activity') / 1000)
|
||||
])
|
||||
drop_categories_columns = fix_categories_values.drop([
|
||||
'id', 'name'
|
||||
])
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to transform the categories DataFrame: {e}")
|
||||
return
|
||||
|
||||
# Write the DataFrame to a new parquet file
|
||||
logging.info("Writing the transformed categories DataFrame to parquet file")
|
||||
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:
|
||||
logging.error(f"Failed to write the transformed categories DataFrame to parquet file: {e}")
|
||||
return
|
||||
@@ -128,15 +138,14 @@ class DimPayees(Dimensions):
|
||||
self.transform()
|
||||
|
||||
def transform(self):
|
||||
# Read the parquet file into a polars DataFrame
|
||||
try:
|
||||
payees_df = pl.read_parquet(self.file_path)
|
||||
source_payees = pl.read_parquet(self.file_path)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to read the base payees parquet file: {e}")
|
||||
return
|
||||
logging.info("Transforming the payees DataFrame")
|
||||
try:
|
||||
payees_df = payees_df.select([
|
||||
base_payees = source_payees.select([
|
||||
'id',
|
||||
'name',
|
||||
'deleted'
|
||||
@@ -144,10 +153,15 @@ class DimPayees(Dimensions):
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to select columns from the payees DataFrame: {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
# Rename the columns
|
||||
payees_df = payees_df.with_columns(pl.col('id').alias('payee_id'))
|
||||
payees_df = payees_df.with_columns(pl.col('name').alias('payee_name'))
|
||||
add_payees_prefix = base_payees.with_columns([
|
||||
pl.col('id').alias('payee_id'),
|
||||
pl.col('name').alias('payee_name')
|
||||
])
|
||||
drop_payees_columns = add_payees_prefix.drop([
|
||||
'id', 'name'
|
||||
])
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to rename columns in the payees DataFrame: {e}")
|
||||
return
|
||||
@@ -155,7 +169,7 @@ class DimPayees(Dimensions):
|
||||
# Write the DataFrame to a new parquet file
|
||||
logging.info("Writing the transformed payees DataFrame to parquet file")
|
||||
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:
|
||||
logging.error(f"Failed to write the transformed payees DataFrame to parquet file: {e}")
|
||||
return
|
||||
@@ -186,11 +200,23 @@ class DimDate(Dimensions):
|
||||
try:
|
||||
# Create a new column to indicate if the date is a weekday or weekend
|
||||
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:
|
||||
logging.error(f"Failed to create a new column to indicate if the date is a weekday or weekend: {e}")
|
||||
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
|
||||
logging.info("Writing the transformed dates DataFrame to parquet file")
|
||||
try:
|
||||
|
||||
+93
-59
@@ -18,49 +18,68 @@ class FactTransactions(Facts):
|
||||
self.transform()
|
||||
|
||||
def transform(self):
|
||||
# Read the parquet file into a polars DataFrame
|
||||
try:
|
||||
transactions_df = pl.read_parquet(self.file_path)
|
||||
source_transactions = pl.read_parquet(self.file_path)
|
||||
except FileNotFoundError:
|
||||
logging.error("The transactions DataFrame does not exist")
|
||||
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")
|
||||
try:
|
||||
transactions_df = (
|
||||
transactions_df
|
||||
.with_columns([
|
||||
pl.col("id").alias("transaction_id"),
|
||||
pl.col("date").alias("transaction_date"),
|
||||
pl.col("amount").alias("transaction_amount"),
|
||||
pl.col("memo").alias("transaction_memo"),
|
||||
pl.col("cleared").alias("transaction_cleared"),
|
||||
pl.col("approved").alias("transaction_approved"),
|
||||
pl.col("flag_color").alias("transaction_flag_color"),
|
||||
pl.col("account_id").alias("account_id"),
|
||||
pl.col("payee_id").alias("payee_id"),
|
||||
pl.col("category_id").alias("category_id"),
|
||||
pl.col("transfer_account_id").alias("transfer_account_id"),
|
||||
])
|
||||
.with_columns([
|
||||
pl.col("memo").fill_null("unknown"),
|
||||
(pl.col("amount") / 100).alias("transaction_amount"),
|
||||
])
|
||||
.drop([
|
||||
"transfer_transaction_id", "matched_transaction_id", "import_id",
|
||||
"subtransactions", "deleted","flag_name","account_name",
|
||||
"payee_name","category_name","import_payee_name","import_payee_name_original",
|
||||
"debt_transaction_type","ingestion_date"
|
||||
])
|
||||
)
|
||||
resolve_transaction_dates = base_transactions.with_columns([
|
||||
pl.col("date").str.strptime(pl.Date, format="%Y-%m-%d").alias("date")
|
||||
])
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to covert the date to date format: {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
add_transaction_prefix = resolve_transaction_dates.with_columns([
|
||||
pl.col("id").alias("transaction_id"),
|
||||
(pl.col("date").dt.year().cast(pl.Utf8) +
|
||||
pl.col("date").dt.month().cast(pl.Utf8).str.zfill(2) +
|
||||
pl.col("date").dt.day().cast(pl.Utf8).str.zfill(2)).alias("transaction_date"),
|
||||
])
|
||||
fix_transaction_nulls = add_transaction_prefix.with_columns([
|
||||
pl.col("memo").fill_null("none"),
|
||||
pl.col("flag_color").fill_null("none"),
|
||||
pl.col("transfer_account_id").fill_null("none"),
|
||||
pl.col("category_id").fill_null("none"),
|
||||
])
|
||||
fix_transaction_values = fix_transaction_nulls.with_columns([
|
||||
(pl.col("amount") / 1000).alias("transaction_amount")
|
||||
])
|
||||
drop_transaction_columns = fix_transaction_values.drop([
|
||||
"id", "date", "amount"
|
||||
])
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to transform the transactions DataFrame: {e}")
|
||||
return
|
||||
# Write the DataFrame to a new parquet file
|
||||
logging.info("Writing the transformed transactions DataFrame to parquet file")
|
||||
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:
|
||||
logging.error(f"Failed to write the transformed transactions DataFrame: {e}")
|
||||
|
||||
@@ -71,46 +90,61 @@ class FactScheduledTransactions(Facts):
|
||||
self.transform()
|
||||
|
||||
def transform(self):
|
||||
# Read the parquet file into a polars DataFrame
|
||||
try:
|
||||
scheduled_transactions_df = pl.read_parquet(self.file_path)
|
||||
source_scheduled = pl.read_parquet(self.file_path)
|
||||
except FileNotFoundError:
|
||||
logging.error("The scheduled transactions DataFrame does not exist")
|
||||
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")
|
||||
try:
|
||||
scheduled_transactions_df = (
|
||||
scheduled_transactions_df
|
||||
.with_columns([
|
||||
pl.col("id").alias("scheduled_transaction_id"),
|
||||
pl.col("date_first").alias("scheduled_transaction_first_date"),
|
||||
pl.col("date_next").alias("scheduled_transaction_next_date"),
|
||||
pl.col("frequency").alias("scheduled_transaction_frequency"),
|
||||
pl.col("amount").alias("scheduled_transaction_amount"),
|
||||
pl.col("memo").alias("scheduled_transaction_memo"),
|
||||
pl.col("flag_color").alias("scheduled_transaction_flag_color"),
|
||||
pl.col("account_id").alias("account_id"),
|
||||
pl.col("payee_id").alias("payee_id"),
|
||||
pl.col("category_id").alias("category_id"),
|
||||
pl.col("transfer_account_id").alias("transfer_account_id"),
|
||||
])
|
||||
.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"
|
||||
])
|
||||
)
|
||||
add_scheduled_prefix = resolve_scheduled_dates.with_columns([
|
||||
pl.col("id").alias("scheduled_transaction_id")
|
||||
])
|
||||
fix_sheduled_nulls = add_scheduled_prefix.with_columns([
|
||||
pl.col("memo").fill_null("none"),
|
||||
pl.col("flag_color").fill_null("none"),
|
||||
pl.col("transfer_account_id").fill_null("none"),
|
||||
pl.col("category_id").fill_null("none"),
|
||||
])
|
||||
fix_scheduled_values = fix_sheduled_nulls.with_columns([
|
||||
(pl.col("amount") / 1000).alias("scheduled_transaction_amount"),
|
||||
])
|
||||
drop_scheduled_columns = fix_scheduled_values.drop([
|
||||
"id", "amount"
|
||||
])
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to transform the scheduled transactions DataFrame: {e}")
|
||||
return
|
||||
# Write the DataFrame to a new parquet file
|
||||
logging.info("Writing the transformed scheduled transactions DataFrame to parquet file")
|
||||
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:
|
||||
logging.error(f"Failed to write the transformed scheduled transactions DataFrame: {e}")
|
||||
|
||||
+26
-38
@@ -4,13 +4,11 @@ import json
|
||||
import logging
|
||||
import requests
|
||||
import sys
|
||||
import yaml
|
||||
from typing import Dict, Any
|
||||
import config.exit_codes as ec
|
||||
|
||||
class Ingest:
|
||||
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
"""
|
||||
Initialize the Ingest class with the provided configuration.
|
||||
@@ -22,19 +20,9 @@ class Ingest:
|
||||
self.entities = config['entities']
|
||||
self.raw_data_path = config['raw_data_path']
|
||||
self.headers = {'Authorization': f'Bearer {self.api_token}'}
|
||||
self.knowledge_cache = self.load_knowledge_cache()
|
||||
self.MAX_RETRIES = config['REQUESTS_MAX_RETRIES']
|
||||
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]):
|
||||
"""
|
||||
@@ -50,8 +38,18 @@ class Ingest:
|
||||
with open(entity_file, 'w') as f:
|
||||
json.dump(data, f, indent=4)
|
||||
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):
|
||||
"""
|
||||
@@ -70,24 +68,14 @@ class Ingest:
|
||||
with open(self.knowledge_file, 'w') as f:
|
||||
json.dump(knowledge_cache, f, indent=4)
|
||||
|
||||
def check_rate_limit(self, response: requests.Response):
|
||||
"""
|
||||
Check and handle the rate limit based on the response headers.
|
||||
"""
|
||||
rate_limit_header = response.headers.get('X-Rate-Limit')
|
||||
if rate_limit_header:
|
||||
requests_made, limit = map(int, rate_limit_header.split('/'))
|
||||
remaining_requests = limit - requests_made
|
||||
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.")
|
||||
knowledge_cache = self.load_knowledge_cache()
|
||||
knowledge_cache[entity] = server_knowledge
|
||||
try:
|
||||
with open(self.knowledge_file, 'w') as f:
|
||||
json.dump(knowledge_cache, f, indent=4)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to update knowledge cache for {entity} in {self.knowledge_file}")
|
||||
raise e
|
||||
|
||||
def handle_response(self, response) -> bool:
|
||||
if response.status_code == 400:
|
||||
@@ -100,7 +88,7 @@ class Ingest:
|
||||
logging.error("Forbidden. Access is denied.")
|
||||
sys.exit(ec.FORBIDDEN)
|
||||
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)
|
||||
elif response.status_code == 409:
|
||||
logging.error("Conflict. The resource cannot be saved due to a conflict.")
|
||||
@@ -118,7 +106,7 @@ class Ingest:
|
||||
response.raise_for_status()
|
||||
return False
|
||||
|
||||
def fetch_and_cache_entity_data(self):
|
||||
def start_ingestion(self):
|
||||
"""
|
||||
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.")
|
||||
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.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}'
|
||||
|
||||
response = None
|
||||
for attempt in range(self.MAX_RETRIES):
|
||||
try:
|
||||
response = requests.get(url, headers=self.headers)
|
||||
@@ -148,6 +138,7 @@ class Ingest:
|
||||
sys.exit(ec.REQUESTS_ERROR)
|
||||
|
||||
data = response.json()
|
||||
logging.debug(f'response data: {data}')
|
||||
server_knowledge = data['data'].get('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)
|
||||
else:
|
||||
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.
|
||||
|
||||
@@ -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')
|
||||
@@ -68,6 +68,12 @@ Then move the files back in one at a time oldest to newest and run again for eac
|
||||
return False
|
||||
|
||||
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)
|
||||
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(
|
||||
pl.when(pl.col(col).is_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)
|
||||
)
|
||||
return df
|
||||
|
||||
@@ -2,3 +2,10 @@ python-dotenv
|
||||
polars
|
||||
requests
|
||||
pyyaml
|
||||
#visualisation requirements below
|
||||
dash
|
||||
pandas
|
||||
pyarrow
|
||||
dash-bootstrap-components
|
||||
# testing requirements below
|
||||
pytest
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user