
from qutip.nonmarkov.heom import (
    BathExponent,
    Bath,
    BosonicBath,
    DrudeLorentzBath,
    DrudeLorentzPadeBath,
    UnderDampedBath,
    FermionicBath,
    LorentzianBath,
    LorentzianPadeBath,
    HEOMSolver,
    HSolverDL,
)


class TestBathAPI:
    def test_api(self):
        # just assert that the baths are importable
        assert BathExponent
        assert Bath
        assert BosonicBath
        assert DrudeLorentzBath
        assert DrudeLorentzPadeBath
        assert UnderDampedBath
        assert FermionicBath
        assert LorentzianBath
        assert LorentzianPadeBath


class TestSolverAPI:
    def test_api(self):
        # just assert that the solvers are importable
        assert HEOMSolver
        assert HSolverDL


from qutip.nonmarkov.heom import (
    BathExponent,
    Bath,
    BosonicBath,
    DrudeLorentzBath,
    DrudeLorentzPadeBath,
    UnderDampedBath,
    FermionicBath,
    LorentzianBath,
    LorentzianPadeBath,
    HEOMSolver,
    HSolverDL,
    HierarchyADOs,
    HierarchyADOsState,
)


class TestBathAPI:
    def test_api(self):
        # just assert that the baths are importable
        assert BathExponent
        assert Bath
        assert BosonicBath
        assert DrudeLorentzBath
        assert DrudeLorentzPadeBath
        assert UnderDampedBath
        assert FermionicBath
        assert LorentzianBath
        assert LorentzianPadeBath


class TestSolverAPI:
    def test_api(self):
        # just assert that the solvers and associated classes are importable
        assert HEOMSolver
        assert HSolverDL
        assert HierarchyADOs
        assert HierarchyADOsState

Test that HierarchyADOs and HierarchyADOsState are part of the HEOM api.
Test that HierarchyADOs and HierarchyADOsState are part of the HEOM api.

  
  from qutip.nonmarkov.heom import (
      BathExponent,
      Bath,
      BosonicBath,
      DrudeLorentzBath,
      DrudeLorentzPadeBath,
      UnderDampedBath,
      FermionicBath,
      LorentzianBath,
      LorentzianPadeBath,
      HEOMSolver,
      HSolverDL,
+     HierarchyADOs,
+     HierarchyADOsState,
  )
  
  
  class TestBathAPI:
      def test_api(self):
          # just assert that the baths are importable
          assert BathExponent
          assert Bath
          assert BosonicBath
          assert DrudeLorentzBath
          assert DrudeLorentzPadeBath
          assert UnderDampedBath
          assert FermionicBath
          assert LorentzianBath
          assert LorentzianPadeBath
  
  
  class TestSolverAPI:
      def test_api(self):
-         # just assert that the solvers are importable
+         # just assert that the solvers and associated classes are importable
          assert HEOMSolver
          assert HSolverDL
+         assert HierarchyADOs
+         assert HierarchyADOsState
  
Test that HierarchyADOs and HierarchyADOsState are part of the HEOM api.
## Code Before:

from qutip.nonmarkov.heom import (
    BathExponent,
    Bath,
    BosonicBath,
    DrudeLorentzBath,
    DrudeLorentzPadeBath,
    UnderDampedBath,
    FermionicBath,
    LorentzianBath,
    LorentzianPadeBath,
    HEOMSolver,
    HSolverDL,
)


class TestBathAPI:
    def test_api(self):
        # just assert that the baths are importable
        assert BathExponent
        assert Bath
        assert BosonicBath
        assert DrudeLorentzBath
        assert DrudeLorentzPadeBath
        assert UnderDampedBath
        assert FermionicBath
        assert LorentzianBath
        assert LorentzianPadeBath


class TestSolverAPI:
    def test_api(self):
        # just assert that the solvers are importable
        assert HEOMSolver
        assert HSolverDL

## Instruction:
Test that HierarchyADOs and HierarchyADOsState are part of the HEOM api.
## Code After:

from qutip.nonmarkov.heom import (
    BathExponent,
    Bath,
    BosonicBath,
    DrudeLorentzBath,
    DrudeLorentzPadeBath,
    UnderDampedBath,
    FermionicBath,
    LorentzianBath,
    LorentzianPadeBath,
    HEOMSolver,
    HSolverDL,
    HierarchyADOs,
    HierarchyADOsState,
)


class TestBathAPI:
    def test_api(self):
        # just assert that the baths are importable
        assert BathExponent
        assert Bath
        assert BosonicBath
        assert DrudeLorentzBath
        assert DrudeLorentzPadeBath
        assert UnderDampedBath
        assert FermionicBath
        assert LorentzianBath
        assert LorentzianPadeBath


class TestSolverAPI:
    def test_api(self):
        # just assert that the solvers and associated classes are importable
        assert HEOMSolver
        assert HSolverDL
        assert HierarchyADOs
        assert HierarchyADOsState
