Skip to content

Commit

Permalink
Add DatabentoDataLoader tests
Browse files Browse the repository at this point in the history
  • Loading branch information
cjdsellers committed Nov 4, 2023
1 parent adbdf44 commit 62cd635
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 1 deletion.
58 changes: 57 additions & 1 deletion nautilus_trader/adapters/databento/loaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,10 @@ class DatabentoDataLoader:
- OHLCV_1D
- DEFINITION
For the loader to work correctly, you must first either:
- Load Databento instrument definitions from a DBN file using `load_instruments(...)`
- Manually add Nautilus instrument objects through `add_instruments(...)`
"""

def __init__(self) -> None:
Expand All @@ -84,6 +88,36 @@ def __init__(self) -> None:

self.load_publishers(path=Path(__file__).resolve().parent / "publishers.json")

def publishers(self) -> dict[int, DatabentoPublisher]:
"""
Return the internal Databento publishers currently held by the loader.
Returns
-------
dict[int, DatabentoPublisher]
Notes
-----
Returns a copy of the internal dictionary.
"""
return self._publishers.copy()

def instruments(self) -> dict[InstrumentId, Instrument]:
"""
Return the internal Nautilus instruments currently held by the loader.
Returns
-------
dict[InstrumentId, Instrument]
Notes
-----
Returns a copy of the internal dictionary.
"""
return self._instruments.copy()

def load_publishers(self, path: PathLike[str] | str) -> None:
"""
Load publisher details from the JSON file at the given path.
Expand Down Expand Up @@ -118,7 +152,29 @@ def load_instruments(self, path: PathLike[str] | str) -> None:
# TODO: Validate actually definitions schema
instruments = self.from_dbn(path)

self._instruments = {i.instrument_id: i for i in instruments}
self._instruments = {i.id: i for i in instruments}

def add_instruments(self, instrument: Instrument | list[Instrument]) -> None:
"""
Add the given `instrument`(s) for use by the loader.
Parameters
----------
instrument : Instrument | list[Instrument]
The Nautilus instrument(s) to add.
Warnings
--------
Will overwrite any existing instrument(s) with the same Nautilus instrument ID(s).
"""
if not isinstance(instrument, list):
instruments = [instrument]
else:
instruments = instrument

for inst in instruments:
self._instruments[inst.id] = inst

def from_dbn(self, path: PathLike[str] | str) -> list[Data]:
"""
Expand Down
37 changes: 37 additions & 0 deletions tests/integration_tests/adapters/databento/test_loaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,43 @@ def test_loader_definition_glbx_all_symbols() -> None:
assert len(data) == 10_000_000


def test_get_publishers() -> None:
# Arrange
loader = DatabentoDataLoader()

# Act
result = loader.publishers()

# Assert
assert len(result) == 43 # From built-in map


def test_get_instruments_when_no_instruments() -> None:
# Arrange
loader = DatabentoDataLoader()

# Act
result = loader.instruments()

# Assert
assert len(result) == 0 # No instruments loaded yet


def test_get_instruments() -> None:
# Arrange
loader = DatabentoDataLoader()
path = DATABENTO_TEST_DATA_DIR / "definition-glbx-es-fut.dbn.zst"

# Act
loader.load_instruments(path)
instruments = loader.instruments()

# Assert
expected_id = InstrumentId.from_str("ESM3.GLBX")
assert len(instruments) == 1
assert instruments[expected_id].id == expected_id


def test_loader_definition_glbx_futures() -> None:
# Arrange
loader = DatabentoDataLoader()
Expand Down

0 comments on commit 62cd635

Please sign in to comment.