テスト用のヘルパークラスで OpenSearch を用いたコードのテストをやりやすくする

OpenSearch や Elasticsearch を使ったコードをテストする際、テスト用のインデックスを作成したり、削除したりする必要があります。
テスト用にコンテナを別途立ち上げても良いのですが、追加でコンテナ立ち上げたりするのがいやだったため、試しに今回は開発環境で使用しているコンテナを使い回すことを前提として、テスト用のインデックスを作成、削除するヘルパークラスを作成します。

今回は OpenSearch を使っていますが、Elasticsearch でも同様に使えると思います。

helper.py

class OpenSearchTestHelper:
    def __init__(self, hosts=[{"host": "localhost", "port": 9200}]):
        user = "user"
        password = "password"
        
        auth = (user, password)
        self.client = OpenSearch(
            hosts=hosts,
            http_compress=True,
            http_auth=auth,
            use_ssl=True,
            verify_certs=False,
            ssl_show_warn=False,
        )
        self.test_indices = set()

    def create_index(self, index_name: str):
        """
        テスト用のインデックスを作成し、クリーンアップリストに追加
        
        Args:
            index_name (str): インデックス名
            
        """
        self.client.indices.create(index=index_name, body={}) # 必要に応じて設定を追加
        self.test_indices.add(index_name)

    def cleanup_indices(self, additional_indices: list[str] | None = None):
        """
        登録された全てのテストインデックスを削除
        
        Args:
            additional_indices (list[str]): 追加で削除するインデックス名。テスト対象コード内で作成したインデックスを削除する際に使用
        """
        if additional_indices is None:
            additional_indices = []

        for index in additional_indices:
            if self.client.indices.exists(index=index):
                self.client.indices.delete(index=index)

        for index in self.test_indices:
            if self.client.indices.exists(index=index):
                self.client.indices.delete(index=index)
                
        self.test_indices.clear()

    def refresh(self, index_name: str):
        """
        指定したインデックスをリフレッシュ
        
        Args:
            index_name (str): インデックス名
        """
        self.client.indices.refresh(index=index_name)

pytestを使っている場合、 conftest.py にヘルパークラスを使うための fixture を追加します。

conftest.py

@pytest.fixture
def opensearch_helper():
    helper = OpenSearchTestHelper()
    yield helper
    helper.cleanup_indices()

これで、テスト用のインデックスを作成、削除するヘルパークラスを使って、テストコードを書くことができます。
テストが終了した際には、作成したインデックスが削除されるため、テスト後のクリーンアップが不要になります。

test_code.py

def test_create_index(opensearch_helper):
    index_name = opensearch_helper.create_index("test_index")
    
    sut = SomeClass()
    sut.some_method()
    
    assert not opensearch_helper.client.indices.exists(index=index_name)

/以上