Bug fix no more request limit (#18)

* added tests

* Removed method no longer used due to YNAB api Changes

.gitignore

@@ -7,4 +7,6 @@ data/*
 __pycache__/*
 */__pycache__/*
 *.pbix
-/logs/*
+/logs/*
+.vscode/*
+*.coverage

config/config.yaml

@@ -25,4 +25,4 @@ processed_data_path: data/processed
 base_data_path: data/base
 warehouse_data_path: data/warehouse
 REQUESTS_MAX_RETRIES: 3
-REQUESTS_RETRY_DELAY: 5
+REQUESTS_RETRY_DELAY: 5

config/exit_codes.py

@@ -10,4 +10,7 @@ NOT_FOUND = 8
 CONFLICT = 9
 MOVE_FILE_ERROR = 10
 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

dash_app.py

@@ -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
+)

docs/ERD.md

@@ -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"
 ```
-

docs/Get_Started.md

@@ -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

docs/dataflow.md

@@ -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.

main.py

@@ -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:

pipeline/dimensions.py

@@ -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,29 +105,32 @@ 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
-        
+
 class DimPayees(Dimensions):
     def __init__(self, config):
         super().__init__(config)
@@ -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:

pipeline/facts.py

@@ -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}")

pipeline/ingest.py

@@ -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):
         """
@@ -64,31 +62,21 @@ class Ingest:
             logging.info(f"Knowledge file not found. Creating a new one at {self.knowledge_file}. This is normal for the first run.")
             os.makedirs(os.path.dirname(self.knowledge_file), exist_ok=True)
             knowledge_cache = {}
-        
+
         knowledge_cache[entity] = server_knowledge
-        
+
         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:
             logging.error("Bad request. The request could not be understood by the API due to malformed syntax or validation errors.")
@@ -100,14 +88,14 @@ 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.")
             sys.exit(ec.CONFLICT)
         elif response.status_code == 429:
             logging.error("Too many requests. You have made too many requests in a short amount of time.")
-            return True 
+            return True
         elif response.status_code == 500:
             logging.error("Internal server error. The API experienced an unexpected error.")
             return True
@@ -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)
@@ -146,11 +136,12 @@ class Ingest:
                     else:
                         logging.error("Max retries reached. Exiting.")
                         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}')
-            
+
             if server_knowledge is not None and server_knowledge != last_knowledge:
                 self.update_server_knowledge_cache(entity, server_knowledge)
                 entity_data = data['data']
@@ -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.

pipeline/pipeline_main.py

@@ -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')

pipeline/raw_to_base.py

@@ -39,20 +39,20 @@ class RawToBase:
                 logging.error(f"entity: {entity} has been processed, but we could not move the file out of the raw folder, please clear the raw folder for {entity}.")
                 sys.exit(ec.MOVE_FILE_ERROR)
             logging.info(f"Successfully processed entity: {entity}")
-    
+
     def _load_raw_data(self, entity):
         entity_path = os.path.join(self.raw_data_path, entity)
         self.data[entity] = []
         logging.debug(f"Loading data for entity: {entity} from path: {entity_path}")
-        
+
         files = [f for f in os.listdir(entity_path) if f.endswith('.json')]
-        
+
         if len(files) > 1:
             logging.error(f"""More than one file found in path: {entity_path}. Skipping processing for entity: {entity}.
 recommended actions is to move the newest file(s) out, re-run main.py.
 Then move the files back in one at a time oldest to newest and run again for each file""")
             return False
-        
+
         if len(files) == 1:
             file_name = files[0]
             file_path = os.path.join(entity_path, file_name)
@@ -63,11 +63,17 @@ Then move the files back in one at a time oldest to newest and run again for eac
             except Exception as e:
                 logging.error(f"Failed to load data from file: {file_path}, error: {e}")
                 return False
-            
+
             if self._is_data_empty(entity, data, file_path):
                 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}")
@@ -88,11 +94,11 @@ Then move the files back in one at a time oldest to newest and run again for eac
                 return True
         logging.debug(f"Data is not empty for entity: {entity}")
         return False
-    
+
     def _add_ingestion_date(self, entity, data, file_name):
         modified_data = []
         ingestion_date = datetime.strptime(file_name.split('.')[0], '%Y%m%d%H%M%S').date()
-        
+
         logging.debug(f"Adding ingestion date to data for entity: {entity}")
         if entity == 'categories':
             for group in data.get('category_groups', []):
@@ -122,7 +128,7 @@ Then move the files back in one at a time oldest to newest and run again for eac
         else:
             self.base_data[entity] = pl.DataFrame()
             logging.debug(f"No existing base data found for entity: {entity}, starting with an empty DataFrame")
-    
+
     #Function to cast null Struct({'': Null}) columns to String
     def _cast_struct_to_string(self,df):
         for col in df.columns:
@@ -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
@@ -138,7 +144,7 @@ Then move the files back in one at a time oldest to newest and run again for eac
     def _combine_data(self, entity):
         logging.debug(f"Combining data for entity: {entity}")
         combined_data = []
-        
+
         # Combine data from the entity
         if entity == 'categories':
             for data in self.data[entity]:
@@ -147,41 +153,41 @@ Then move the files back in one at a time oldest to newest and run again for eac
         else:
             for data in self.data[entity]:
                 combined_data.extend(data)
-        
+
         new_data_df = pl.DataFrame(combined_data)
-        
+
         # Ensure the unique id column is preserved
         unique_id = self.primary_keys[entity]['unique_id']
         if unique_id not in new_data_df.columns:
             logging.error(f"Unique ID column '{unique_id}' not found in the combined data for entity: {entity}")
             exit(ec.UNIQUE_ID_NOT_FOUND)
-        
+
         # Cast columns in new_data_df
         new_data_df = self._cast_struct_to_string(new_data_df)
-        
+
         # Merge new data with existing base data
         if entity in self.base_data and not self.base_data[entity].is_empty():
             existing_data_df = self.base_data[entity]
-            
+
             # Cast columns in existing_data_df
             existing_data_df = self._cast_struct_to_string(existing_data_df)
-            
+
             # Identify new rows and rows to update
             new_rows = new_data_df.filter(~pl.col(unique_id).is_in(existing_data_df[unique_id]))
             updated_rows = new_data_df.filter(pl.col(unique_id).is_in(existing_data_df[unique_id]))
-            
+
             # Update existing rows
             for row in updated_rows.iter_rows(named=True):
                 existing_data_df = existing_data_df.with_columns([
                     pl.when(pl.col(unique_id) == row[unique_id]).then(pl.lit(row[col], allow_object=True)).otherwise(pl.col(col)).alias(col)
                     for col in updated_rows.columns if col != unique_id
                 ])
-            
+
             # Add new rows
             self.base_data[entity] = pl.concat([existing_data_df, new_rows])
         else:
             self.base_data[entity] = new_data_df
-        
+
         logging.debug(f"Successfully combined data for entity: {entity}")
 
     def _save_base_data(self, entity):
@@ -194,34 +200,34 @@ Then move the files back in one at a time oldest to newest and run again for eac
             return False
         logging.debug(f"Saved base data for entity: {entity} to path: {file_path}")
         return True
-    
+
     def _move_raw_to_processed(self, entity):
         raw_entity_path = os.path.join(self.raw_data_path, entity)
         processed_path = os.path.join(self.processed_data_path, entity)
-        
+
         os.makedirs(processed_path, exist_ok=True)
-        
+
         try:
             files = [f for f in os.listdir(raw_entity_path) if f.endswith('.json')]
             if len(files) != 1:
                 logging.error(f"Expected exactly one file in path: {raw_entity_path}, but found {len(files)}")
                 return False
-            
+
             file_name = files[0]
             raw_file_path = os.path.join(raw_entity_path, file_name)
             processed_file_path = os.path.join(processed_path, file_name)
-            
+
             logging.debug(f"Moving file: {raw_file_path} to {processed_file_path}")
-            
+
             os.rename(raw_file_path, processed_file_path)
             logging.debug(f"Moved file: {file_name} to processed")
-        
+
         except FileNotFoundError as e:
             logging.error(f"File not found: {e}")
             return False
         except Exception as e:
             logging.error(f"Failed to move file for entity: {entity}, error: {e}")
             return False
-        
+
         logging.debug(f"Moved processed file for entity: {entity} to path: {processed_path}")
-        return True
+        return True

requirements.txt

@@ -1,4 +1,11 @@
 python-dotenv
 polars
 requests
-pyyaml
+pyyaml
+#visualisation requirements below
+dash
+pandas
+pyarrow 
+dash-bootstrap-components
+# testing requirements below
+pytest

tests/test_ingest.py

@@ -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()

tests/test_main.py

@@ -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)