11import unittest
22from unittest .mock import patch , MagicMock
3- from yfinance .const import PREDEFINED_SCREENER_BODY_MAP
4- from yfinance .screener .screener import Screener
5- from yfinance .screener .screener_query import EquityQuery
3+ from yfinance .screener .screener import screen
4+ from yfinance .screener .query import EquityQuery
65
76
87class TestScreener (unittest .TestCase ):
98
109 @classmethod
1110 def setUpClass (self ):
12- self .screener = Screener ()
1311 self .query = EquityQuery ('gt' ,['eodprice' ,3 ])
14-
15- def test_set_default_body (self ):
16- result = self .screener .set_default_body (self .query )
17-
18- self .assertEqual (self .screener .body ['offset' ], 0 )
19- self .assertEqual (self .screener .body ['size' ], 100 )
20- self .assertEqual (self .screener .body ['sortField' ], 'ticker' )
21- self .assertEqual (self .screener .body ['sortType' ], 'desc' )
22- self .assertEqual (self .screener .body ['quoteType' ], 'equity' )
23- self .assertEqual (self .screener .body ['query' ], self .query .to_dict ())
24- self .assertEqual (self .screener .body ['userId' ], '' )
25- self .assertEqual (self .screener .body ['userIdType' ], 'guid' )
26- self .assertEqual (self .screener , result )
27-
28- def test_set_predefined_body (self ):
29- k = 'most_actives'
30- result = self .screener .set_predefined_body (k )
31- self .assertEqual (self .screener .body , PREDEFINED_SCREENER_BODY_MAP [k ])
32- self .assertEqual (self .screener , result )
33-
34- def test_set_predefined_body_invalid_key (self ):
35- with self .assertRaises (ValueError ):
36- self .screener .set_predefined_body ('invalid_key' )
37-
38- def test_set_body (self ):
39- body = {
40- "offset" : 0 ,
41- "size" : 100 ,
42- "sortField" : "ticker" ,
43- "sortType" : "desc" ,
44- "quoteType" : "equity" ,
45- "query" : self .query .to_dict (),
46- "userId" : "" ,
47- "userIdType" : "guid"
48- }
49- result = self .screener .set_body (body )
50-
51- self .assertEqual (self .screener .body , body )
52- self .assertEqual (self .screener , result )
53-
54- def test_set_body_missing_keys (self ):
55- body = {
56- "offset" : 0 ,
57- "size" : 100 ,
58- "sortField" : "ticker" ,
59- "sortType" : "desc" ,
60- "quoteType" : "equity"
61- }
62- with self .assertRaises (ValueError ):
63- self .screener .set_body (body )
64-
65- def test_set_body_extra_keys (self ):
66- body = {
67- "offset" : 0 ,
68- "size" : 100 ,
69- "sortField" : "ticker" ,
70- "sortType" : "desc" ,
71- "quoteType" : "equity" ,
72- "query" : self .query .to_dict (),
73- "userId" : "" ,
74- "userIdType" : "guid" ,
75- "extraKey" : "extraValue"
76- }
77- with self .assertRaises (ValueError ):
78- self .screener .set_body (body )
79-
80- def test_patch_body (self ):
81- initial_body = {
82- "offset" : 0 ,
83- "size" : 100 ,
84- "sortField" : "ticker" ,
85- "sortType" : "desc" ,
86- "quoteType" : "equity" ,
87- "query" : self .query .to_dict (),
88- "userId" : "" ,
89- "userIdType" : "guid"
90- }
91- self .screener .set_body (initial_body )
92- patch_values = {"size" : 50 }
93- result = self .screener .patch_body (patch_values )
94-
95- self .assertEqual (self .screener .body ['size' ], 50 )
96- self .assertEqual (self .screener .body ['query' ], self .query .to_dict ())
97- self .assertEqual (self .screener , result )
98-
99- def test_patch_body_extra_keys (self ):
100- initial_body = {
101- "offset" : 0 ,
102- "size" : 100 ,
103- "sortField" : "ticker" ,
104- "sortType" : "desc" ,
105- "quoteType" : "equity" ,
106- "query" : self .query .to_dict (),
107- "userId" : "" ,
108- "userIdType" : "guid"
109- }
110- self .screener .set_body (initial_body )
111- patch_values = {"extraKey" : "extraValue" }
112- with self .assertRaises (ValueError ):
113- self .screener .patch_body (patch_values )
12+ self .predefined = 'aggressive_small_caps'
11413
11514 @patch ('yfinance.screener.screener.YfData.post' )
11615 def test_set_large_size_in_body (self , mock_post ):
117- body = {
118- "offset" : 0 ,
119- "size" : 251 , # yahoo limits at 250
120- "sortField" : "ticker" ,
121- "sortType" : "desc" ,
122- "quoteType" : "equity" ,
123- "query" : self .query .to_dict (),
124- "userId" : "" ,
125- "userIdType" : "guid"
126- }
127-
12816 with self .assertRaises (ValueError ):
129- self .screener . set_body ( body ). response
17+ screen ( self .query , size = 251 )
13018
131- @patch ('yfinance.screener.screener .YfData.post' )
132- def test_fetch (self , mock_post ):
19+ @patch ('yfinance.data .YfData.post' )
20+ def test_fetch_query (self , mock_post ):
13321 mock_response = MagicMock ()
134- mock_response .json .return_value = {'finance' : {'result' : [{}]}}
22+ mock_response .json .return_value = {'finance' : {'result' : [{'key' : 'value' }]}}
13523 mock_post .return_value = mock_response
13624
137- self .screener .set_default_body (self .query )
138- response = self .screener ._fetch ()
139-
140- self .assertEqual (response , {'finance' : {'result' : [{}]}})
25+ response = screen (self .query )
26+ self .assertEqual (response , {'key' : 'value' })
14127
142- @patch ('yfinance.screener.screener. YfData.post ' )
143- def test_fetch_and_parse (self , mock_post ):
28+ @patch ('yfinance.data. YfData.get ' )
29+ def test_fetch_predefined (self , mock_get ):
14430 mock_response = MagicMock ()
14531 mock_response .json .return_value = {'finance' : {'result' : [{'key' : 'value' }]}}
146- mock_post .return_value = mock_response
32+ mock_get .return_value = mock_response
14733
148- self .screener .set_default_body (self .query )
149- self .screener ._fetch_and_parse ()
150- self .assertEqual (self .screener .response , {'key' : 'value' })
34+ response = screen (self .predefined )
35+ self .assertEqual (response , {'key' : 'value' })
15136
15237if __name__ == '__main__' :
15338 unittest .main ()
0 commit comments