Skip to content

Commit

Permalink
tests
Browse files Browse the repository at this point in the history
  • Loading branch information
abeltavares committed May 24, 2024
1 parent f06e828 commit eb03238
Show file tree
Hide file tree
Showing 6 changed files with 316 additions and 219 deletions.
63 changes: 22 additions & 41 deletions tests/dags_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,35 +5,42 @@
from airflow.models import DagBag
from airflow.operators.python import PythonOperator
import logging
from core.market_data_processor import StockApiClient, CryptoApiClient, Storage
from dags.market_data_dag import process_crypto_data_task, process_stock_data_task

# Set the logging level to ERROR for the Airflow logger
logging.getLogger("airflow").setLevel(logging.ERROR)

# Find the parent directory
parent_directory = os.path.dirname(os.path.abspath(__file__))

# Find the project root
project_root = os.path.dirname(parent_directory)

# Add the project root to the Python path
sys.path.insert(0, project_root)

from core.market_data_processor import StockApiClient, CryptoApiClient, Storage
from dags.market_data_dag import process_crypto_data_task, process_stock_data_task


class TestMarketDataDag(unittest.TestCase):
"""
Unit tests for the Market Data DAGs.
"""

def setUp(self):

self.dagbag = DagBag(
dag_folder=os.path.join(project_root, "dags"), include_examples=False
)
self.stock_dag_id = "process_stock_data"
self.crypto_dag_id = "process_crypto_data"
with patch(
"dags.market_data_dag.read_json",
return_value={
"assets": {
"stocks": {"symbols": ["BREPE"], "schedule_interval": "0 13 * * *"},
"cryptos": {
"symbols": ["BTC", "ETH"],
"schedule_interval": "0 13 * * *",
},
}
},
):
self.dagbag = DagBag(
dag_folder=os.path.join(project_root, "dags"), include_examples=False
)
self.stock_dag_id = "process_stock_data"
self.crypto_dag_id = "process_crypto_data"

def test_dag_stocks_exists(self):
self.assertIn(self.stock_dag_id, self.dagbag.dags)
Expand All @@ -46,33 +53,11 @@ def test_dag_stocks_loaded(self):

def test_dag_stocks_schedule_interval(self):
dag = self.dagbag.get_dag(self.stock_dag_id)
self.assertEqual(dag.schedule_interval, "0 23 * * 1-5")
self.assertEqual(dag.schedule_interval, "0 13 * * *")

@patch.object(StockApiClient, "get_stocks")
@patch.object(StockApiClient, "get_data")
@patch.object(Storage, "store_data")
def test_process_stock_data_task(
self, mock_store_data, mock_get_data, mock_get_stocks
):
# Setup mock behavior
stocks = {"gainers": ["ABC"]}

stock_data = {
"gainers": [
{
"symbol": "ABC",
"volume": "123456",
"price": "50.25",
"change_percent": "2.5",
"market_cap": "1.2B",
"name": "ABC Company",
}
]
}
mock_get_stocks.return_value = stocks
mock_get_data.return_value = stock_data

# Get the task
def test_process_stock_data_task(self, mock_store_data, mock_get_data):
task_id = "get_stocks"

test = PythonOperator(
Expand All @@ -82,8 +67,6 @@ def test_process_stock_data_task(

test.execute(context={})

# Check if the methods were called
mock_get_stocks.assert_called_once()
mock_get_data.assert_called_once()
mock_store_data.assert_called_once()

Expand All @@ -98,12 +81,11 @@ def test_dag_cryptos_loaded(self):

def test_dag_cryptos_schedule_interval(self):
dag = self.dagbag.get_dag(self.crypto_dag_id)
self.assertEqual(dag.schedule_interval, "0 23 * * *")
self.assertEqual(dag.schedule_interval, "0 13 * * *")

@patch.object(CryptoApiClient, "get_data")
@patch.object(Storage, "store_data")
def test_process_crypto_data_task(self, mock_get_crypto_data, mock_store_data):
# Get the DAG and task
mock_get_crypto_data.return_value = {}
task_id = "get_crypto"

Expand All @@ -114,7 +96,6 @@ def test_process_crypto_data_task(self, mock_get_crypto_data, mock_store_data):

test.execute(context={})

# Check if the methods were called
mock_get_crypto_data.assert_called_once()
mock_store_data.assert_called_once()

Expand Down
1 change: 1 addition & 0 deletions tests/test_base_api_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
# TODO: Add tests for the BaseApiClient class
73 changes: 73 additions & 0 deletions tests/test_crypto_api_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
import unittest
from unittest.mock import patch, MagicMock
import logging
from utils import market_data_processor_utils
from core.crypto_api_client import CryptoApiClient


class TestCryptoApiClient(unittest.TestCase):
"""
Unit tests for the CryptoApiClient class.
"""

def setUp(self):
self.logger = MagicMock(spec=logging.Logger)
with patch(
"core.crypto_api_client.read_json",
return_value={"assets": {"cryptos": {"symbols": ["BREPE"]}}},
):
self.crypto_api_client = CryptoApiClient(logger=self.logger)

@patch("core.crypto_api_client.requests.get")
def test_get_data(
self,
mock_get,
):
# Mock the response from the API
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": [
{
"name": "BREPE",
"symbol": "BREPE",
"quote": {
"USD": {
"price": 7,
"volume_24h": 10,
"percent_change_24h": 16,
"market_cap": 0,
}
},
}
]
}
mock_get.return_value = mock_response

# Call the method under test
data = self.crypto_api_client.get_data()

# Assert that the method returned the expected data
expected_data = {
"BREPE": {
"price": 7,
"volume": 10,
"change_percent": 16,
"market_cap": 0,
"name": "BREPE",
}
}
self.assertEqual(data, expected_data)

@patch("core.crypto_api_client.requests.get")
def test_get_data_invalid_data(self, mock_get):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": [{"name": "BREPE", "quote": {}}] # Simulate missing quote data
}
mock_get.return_value = mock_response

with patch.object(self.logger, "error"):
with self.assertRaises(KeyError):
self.crypto_api_client.get_data()
179 changes: 1 addition & 178 deletions tests/test_data_processor.py
Original file line number Diff line number Diff line change
@@ -1,178 +1 @@
import os
import sys
import unittest
from unittest.mock import patch, MagicMock
import logging
from utils import market_data_processor_utils
from core.market_data_processor import (
StockApiClient,
CryptoApiClient,
Storage,
MarketDataEngine,
)

class TestStorage(unittest.TestCase):
"""
Unit tests for the Storage class.
"""

def setUp(self):
self.logger = MagicMock(spec=logging.Logger)

self.storage = Storage(logger=self.logger)

@patch.dict(
"core.market_data_processor.os.environ",
{
"POSTGRES_USER": "user",
"POSTGRES_PASSWORD": "password",
"POSTGRES_DB": "test_db",
"POSTGRES_HOST": "localhost",
"POSTGRES_PORT": "5432",
},
clear=True,
)
@patch("core.market_data_processor.psycopg2.connect")
def test_connect(self, mock_connect):
self.storage._connect()

mock_connect.assert_called_once_with(
host="localhost",
port="5432",
database="test_db",
user="user",
password="password",
)

@patch("core.market_data_processor.psycopg2.connect")
def test_close(self, mock_connect):
mock_conn = mock_connect.return_value
mock_cur = MagicMock()
mock_conn.cursor.return_value = mock_cur

self.storage._connect()
self.storage._close()

mock_cur.close.assert_called_once()
mock_conn.close.assert_called_once()

@patch("core.market_data_processor.psycopg2.connect")
def test_store_data_with_valid_data(self, mock_connect):
mock_conn = mock_connect.return_value
mock_cur = MagicMock()
mock_conn.cursor.return_value = mock_cur
mock_execute = mock_cur.execute
mock_commit = mock_conn.commit

# Test case with valid data
data = {
"ABC": {
"volume": 123456,
"price": 50.25,
"change_percent": 2.5,
"market_cap": "1.2B",
"name": "ABC Company",
}
}
table = "stocks"

self.storage.store_data(data, table)

# Assert that execute and commit methods were called
mock_execute.assert_called_once()
mock_execute.assert_called_once_with(
"INSERT INTO stocks (symbol, name, market_cap, volume, price, change_percent) VALUES (%s, %s, %s, %s, %s, %s)",
("ABC", "ABC Company", "1.2B", 123456, 50.25, 2.5),
)
mock_commit.assert_called_once()

@patch("core.market_data_processor.psycopg2.connect")
def test_store_data_with_invalid_data_empty_symbol(self, mock_connect):
# (empty name)
data_invalid = {
"ABC": {
"volume": 1000,
"price": 10.0,
"change_percent": 5.0,
"market_cap": 1000000,
"name": "",
}
}

table = "stocks"

# Mock the logger.error method to capture log messages
with patch.object(self.logger, "error"):
with self.assertRaises(ValueError):
self.storage.store_data(data_invalid, table)

@patch("core.market_data_processor.psycopg2.connect")
def test_store_data_with_invalid_data_type(self, mock_connect):
data = {
"ABC": {
"volume": 1000,
"price": 10.0,
"change_percent": 5.0,
"market_cap": 1000000,
"name": "",
}
}

table_invalid = 123 # Invalid table(not a string)

with patch.object(self.logger, "error"):
with self.assertRaises(TypeError):
self.storage.store_data(data, table_invalid)


class TestMarketDataEngine(unittest.TestCase):
"""
Unit tests for the MarketDataEngine class.
"""

def setUp(self):
self.stock_api_client = MagicMock(spec=StockApiClient)
self.crypto_api_client = MagicMock(spec=CryptoApiClient)
self.db_connector = MagicMock(spec=Storage)
self.logger = MagicMock(spec=logging.Logger)
self.stock_engine = MarketDataEngine(
self.stock_api_client, self.db_connector, self.logger
)
self.crypto_engine = MarketDataEngine(
self.crypto_api_client, self.db_connector, self.logger
)

def test_process_stock_data(self):
self.stock_api_client.get_data.return_value = {
"AAPL": {"price": 150.0},
"GOOG": {"price": 2000.0},
"MSFT": {"price": 300.0},
}

self.stock_engine.process_stock_data()

self.db_connector.store_data.assert_called_once_with(
{
"AAPL": {"price": 150.0},
"GOOG": {"price": 2000.0},
"MSFT": {"price": 300.0},
},
"stocks",
)

def test_process_crypto_data(self):
self.crypto_api_client.get_data.return_value = {
"BTC": {"price": 50000.0},
"ETH": {"price": 4000.0},
}

self.crypto_engine.process_crypto_data()

self.crypto_api_client.get_data.assert_called_once()
self.db_connector.store_data.assert_called_once_with(
{"BTC": {"price": 50000.0}, "ETH": {"price": 4000.0}}, "cryptos"
)


if __name__ == "__main__":
unittest.main()
# TODO: Add tests for the DataProcessor class
Loading

0 comments on commit eb03238

Please sign in to comment.