From 62cd6358b6da67ca1b9fc3aa8be25d47ce5f3544 Mon Sep 17 00:00:00 2001
From: Chris Sellers <chris@cjdsellers.io>
Date: Sat, 4 Nov 2023 17:46:51 +1100
Subject: [PATCH] Add DatabentoDataLoader tests

---
 nautilus_trader/adapters/databento/loaders.py | 58 ++++++++++++++++++-
 .../adapters/databento/test_loaders.py        | 37 ++++++++++++
 2 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/nautilus_trader/adapters/databento/loaders.py b/nautilus_trader/adapters/databento/loaders.py
index d2dd29af8e88..8a9d32a589f7 100644
--- a/nautilus_trader/adapters/databento/loaders.py
+++ b/nautilus_trader/adapters/databento/loaders.py
@@ -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:
@@ -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.
@@ -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]:
         """
diff --git a/tests/integration_tests/adapters/databento/test_loaders.py b/tests/integration_tests/adapters/databento/test_loaders.py
index 03e6483133ba..486aef360551 100644
--- a/tests/integration_tests/adapters/databento/test_loaders.py
+++ b/tests/integration_tests/adapters/databento/test_loaders.py
@@ -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()