Files
OpenBB/examples/openbb-apachebeam/tests/test_obb_pipeline.py
mmistroni a14e0cf9cc Feature/openbb apachebeam (#6679)
* adding beam sample

* adding senate discl

* removingn commits on wrong branch

* amended according to latest comments

* Delete openbb_platform/providers/government_us/openbb_government_us/utils/senate_constants.py

file in another bramnch

* Delete openbb_platform/providers/government_us/openbb_government_us/utils/senate_helpers.py

file belong to antothe rbranch

* Update README.md

fixed comment

* Update README.md

---------

Co-authored-by: Danglewood <85772166+deeleeramone@users.noreply.github.com>
Co-authored-by: Theodore Aptekarev <aptekarev@gmail.com>
2024-09-26 23:20:46 +00:00

49 lines
1.8 KiB
Python

import unittest
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.options.pipeline_options import PipelineOptions
import asyncio
import apache_beam as beam
from openbb_yfinance.models.equity_quote import YFinanceEquityQuoteFetcher as quote_fetcher
from openbb_yfinance.models.equity_profile import YFinanceEquityProfileFetcher as profile_fetcher
from openbb_yfinance.models.company_news import YFinanceCompanyNewsFetcher as news_fetcher
class AsyncProcess(beam.DoFn):
def __init__(self, credentials, fetcher):
self.credentials = credentials
self.fetcher = fetcher
async def fetch_data(self, element: str):
params = dict(symbol=element)
data = await self.fetcher.fetch_data(params, self.credentials)
return [d.model_dump(exclude_none=True) for d in data]
def process(self, element: str):
return asyncio.run(self.fetch_data(element))
class MyTestCase(unittest.TestCase):
def test_sample_pipeline(self):
credentials = {} # Running OBB endpoints which do not require credentials
debug_sink = beam.Map(print)
ticker = 'AAPL'
with TestPipeline(options=PipelineOptions()) as p:
quote = (p | 'Start Quote' >> beam.Create([ticker])
| 'Run Quote' >> beam.ParDo(AsyncProcess(credentials, quote_fetcher))
| 'Print quote' >> debug_sink)
profile = (p | 'Start Profile' >> beam.Create([ticker])
| 'Run Profile' >> beam.ParDo(AsyncProcess(credentials, profile_fetcher))
| 'Print profile' >> debug_sink)
news = (p | 'Start News' >> beam.Create([ticker])
| 'Run News' >> beam.ParDo(AsyncProcess(credentials, news_fetcher))
| 'Print nes' >> debug_sink)
if __name__ == '__main__':
unittest.main()