diff --git a/pyproject.toml b/pyproject.toml index 0fd7d0f..edb35dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "finsim" -version = "2.0.0a1" +version = "2.0.0" authors = [ {name = "Kwan Yuet Stephen Ho", email = "stephenhky@yahoo.com.hk"} ] @@ -56,7 +56,6 @@ packages = [ "finsim.portfolio.optimize", "finsim.portfolio.optimize.native", "finsim.tech", - "finsim.retrieve_stock_symbols_cli", "finsim.schemas" ] zip-safe = false @@ -64,7 +63,7 @@ include-package-data = true package-dir = {"" = "src"} [project.scripts] -retrieve_stock_symbols = "finsim.retrieve_stock_symbols_cli:main_cli" +retrieve_stock_symbols = "finsim.cli:retrieve_stock_symbols" [project.optional-dependencies] test = ["unittest2", "pytest"] diff --git a/src/finsim/retrieve_stock_symbols_cli/cli.py b/src/finsim/cli.py similarity index 89% rename from src/finsim/retrieve_stock_symbols_cli/cli.py rename to src/finsim/cli.py index d6f6a4e..fb2884c 100644 --- a/src/finsim/retrieve_stock_symbols_cli/cli.py +++ b/src/finsim/cli.py @@ -1,8 +1,10 @@ + import os +from typing import Optional import click import pandas as pd -from ..data.finnhub import FinnHubStockReader +from .data.finnhub import FinnHubStockReader @click.command(help='Retrieve stock symbols from Finnhub') @@ -10,7 +12,12 @@ @click.option('--finnhubtokenpath', default=None, help='path of Finnhub tokens') @click.option('--useenvtoken', is_flag=True, default=False, help='Use the environment variable FINNHUBTOKEN as the tokens') @click.option('--shorten', is_flag=True, default=False, help='shorten list of symbols') -def main_cli(outputpath, finnhubtokenpath=None, useenvtoken=False, shorten=False): +def retrieve_stock_symbols( + outputpath: str | os.PathLike, + finnhubtokenpath: Optional[str | os.PathLike] = None, + useenvtoken: bool = False, + shorten: bool = False +) -> None: """Main CLI function to retrieve stock symbols from Finnhub and save them to a file. This function parses command line arguments, retrieves stock symbols from Finnhub, diff --git a/src/finsim/retrieve_stock_symbols_cli/__init__.py b/src/finsim/retrieve_stock_symbols_cli/__init__.py deleted file mode 100644 index b26dbc2..0000000 --- a/src/finsim/retrieve_stock_symbols_cli/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ - -from .cli import main_cli diff --git a/test/test_retrieve_stock_symbols_cli.py b/test/test_retrieve_stock_symbols_cli.py new file mode 100644 index 0000000..f98dbcf --- /dev/null +++ b/test/test_retrieve_stock_symbols_cli.py @@ -0,0 +1,132 @@ +import os +import tempfile +import unittest +from unittest.mock import patch, MagicMock + +from click.testing import CliRunner +import pandas as pd + +from finsim.cli import retrieve_stock_symbols + + +class TestRetrieveStockSymbolsCLI(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + self.sample_symbols = [ + {'symbol': 'AAPL', 'mic': 'XNAS', 'type': 'Common Stock'}, + {'symbol': 'MSFT', 'mic': 'XNAS', 'type': 'Common Stock'}, + {'symbol': 'BRK.A', 'mic': 'XNYS', 'type': 'Common Stock'}, + {'symbol': 'TEST.PUB', 'mic': 'XNAS', 'type': 'PUBLIC'}, + {'symbol': 'OTHER', 'mic': 'BATS', 'type': 'Common Stock'}, + ] + + def test_help(self): + result = self.runner.invoke(retrieve_stock_symbols, ['--help']) + self.assertEqual(result.exit_code, 0) + self.assertIn('Retrieve stock symbols from Finnhub', result.output) + self.assertIn('--finnhubtokenpath', result.output) + self.assertIn('--useenvtoken', result.output) + self.assertIn('--shorten', result.output) + + def test_missing_outputpath(self): + result = self.runner.invoke(retrieve_stock_symbols, []) + self.assertNotEqual(result.exit_code, 0) + self.assertIn('Missing argument', result.output) + + def test_missing_token_options(self): + with tempfile.TemporaryDirectory() as tmpdir: + out_file = os.path.join(tmpdir, 'symbols.json') + result = self.runner.invoke(retrieve_stock_symbols, [out_file]) + self.assertNotEqual(result.exit_code, 0) + self.assertIsInstance(result.exception, ValueError) + self.assertIn('Either --finnhubtokenpath or --useenvtoken must be specified', str(result.exception)) + + def test_useenvtoken_missing_env(self): + with tempfile.TemporaryDirectory() as tmpdir: + out_file = os.path.join(tmpdir, 'symbols.json') + env = os.environ.copy() + env.pop('FINNHUBTOKEN', None) + result = self.runner.invoke(retrieve_stock_symbols, [out_file, '--useenvtoken'], env=env) + self.assertNotEqual(result.exit_code, 0) + self.assertIsInstance(result.exception, ValueError) + self.assertIn('Finnhub tokens not found in the environment variable $FINNHUBTOKEN', str(result.exception)) + + @patch('finsim.cli.FinnHubStockReader') + def test_success_with_useenvtoken_json(self, mock_reader_cls): + mock_reader = MagicMock() + mock_reader.get_all_US_symbols.return_value = self.sample_symbols + mock_reader_cls.return_value = mock_reader + + with tempfile.TemporaryDirectory() as tmpdir: + out_file = os.path.join(tmpdir, 'symbols.json') + result = self.runner.invoke( + retrieve_stock_symbols, + [out_file, '--useenvtoken'], + env={'FINNHUBTOKEN': 'dummy_env_token'} + ) + self.assertEqual(result.exit_code, 0) + mock_reader_cls.assert_called_once_with('dummy_env_token') + self.assertTrue(os.path.exists(out_file)) + + df = pd.read_json(out_file) + self.assertEqual(len(df), len(self.sample_symbols)) + + @patch('finsim.cli.FinnHubStockReader') + def test_success_with_finnhubtokenpath_and_shorten(self, mock_reader_cls): + mock_reader = MagicMock() + mock_reader.get_all_US_symbols.return_value = self.sample_symbols + mock_reader_cls.return_value = mock_reader + + with tempfile.TemporaryDirectory() as tmpdir: + token_file = os.path.join(tmpdir, 'token.txt') + with open(token_file, 'w') as f: + f.write('dummy_file_token\n') + + out_file = os.path.join(tmpdir, 'symbols.csv') + result = self.runner.invoke( + retrieve_stock_symbols, + [out_file, '--finnhubtokenpath', token_file, '--shorten'] + ) + self.assertEqual(result.exit_code, 0) + mock_reader_cls.assert_called_once_with('dummy_file_token') + self.assertTrue(os.path.exists(out_file)) + + df = pd.read_csv(out_file) + # BRK.A has dot, TEST.PUB is PUBLIC, OTHER is BATS -> only AAPL and MSFT remain + symbols = list(df['symbol']) + self.assertEqual(symbols, ['AAPL', 'MSFT']) + + @patch('finsim.cli.FinnHubStockReader') + def test_nonexistent_directory(self, mock_reader_cls): + mock_reader = MagicMock() + mock_reader.get_all_US_symbols.return_value = self.sample_symbols + mock_reader_cls.return_value = mock_reader + + out_file = '/nonexistent/path/for/symbols.json' + result = self.runner.invoke( + retrieve_stock_symbols, + [out_file, '--useenvtoken'], + env={'FINNHUBTOKEN': 'dummy_env_token'} + ) + self.assertNotEqual(result.exit_code, 0) + self.assertIsInstance(result.exception, FileNotFoundError) + + @patch('finsim.cli.FinnHubStockReader') + def test_unrecognized_extension(self, mock_reader_cls): + mock_reader = MagicMock() + mock_reader.get_all_US_symbols.return_value = self.sample_symbols + mock_reader_cls.return_value = mock_reader + + with tempfile.TemporaryDirectory() as tmpdir: + out_file = os.path.join(tmpdir, 'symbols.txt') + result = self.runner.invoke( + retrieve_stock_symbols, + [out_file, '--useenvtoken'], + env={'FINNHUBTOKEN': 'dummy_env_token'} + ) + self.assertNotEqual(result.exit_code, 0) + self.assertIsInstance(result.exception, IOError) + + +if __name__ == '__main__': + unittest.main()