From 9f1351e778bddd667d6a08401a00ad69c6aec3e9 Mon Sep 17 00:00:00 2001
From: Josh Borrow <josh@joshborrow.com>
Date: Wed, 28 Feb 2024 16:32:57 -0500
Subject: [PATCH] Add client ops for store state setting and tests

---
 hera_librarian/client.py                      | 45 +++++++++++++++++++
 tests/integration_test/test_admin_services.py | 33 ++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/hera_librarian/client.py b/hera_librarian/client.py
index 8ae3476..9670d55 100644
--- a/hera_librarian/client.py
+++ b/hera_librarian/client.py
@@ -29,6 +29,8 @@
     AdminStoreListResponse,
     AdminStoreManifestRequest,
     AdminStoreManifestResponse,
+    AdminStoreStateChangeRequest,
+    AdminStoreStateChangeResponse,
 )
 from .models.errors import (
     ErrorClearRequest,
@@ -799,6 +801,49 @@ def get_store_list(
 
         return response.root
 
+    def set_store_state(
+        self,
+        store_name: str,
+        enabled: bool,
+    ) -> bool:
+        """
+        Sets the enabled (or disabled) state of a store on this librarian.
+
+        Parameters
+        ----------
+        store_name : str
+            The name of the store to change the state of.
+        enabled : bool
+            The new state of the store.
+
+        Returns
+        -------
+        bool
+            The new (confirmed) state of the store.
+
+        Raises
+        ------
+        LibrarianError
+            If the store does not exist.
+        """
+
+        try:
+            response: AdminStoreStateChangeResponse = self.post(
+                endpoint="admin/stores/state_change",
+                request=AdminStoreStateChangeRequest(
+                    store_name=store_name,
+                    enabled=enabled,
+                ),
+                response=AdminStoreStateChangeResponse,
+            )
+        except LibrarianHTTPError as e:
+            if e.status_code == 400 and "Store" in e.reason:
+                raise LibrarianError(e.reason)
+            else:
+                raise e
+
+        return response.enabled
+
     def get_store_manifest(
         self,
         store_name: str,
diff --git a/tests/integration_test/test_admin_services.py b/tests/integration_test/test_admin_services.py
index 91da7ba..a6d272f 100644
--- a/tests/integration_test/test_admin_services.py
+++ b/tests/integration_test/test_admin_services.py
@@ -120,3 +120,36 @@ def test_store_manifest(server, admin_client):
             assert entry.instance_available is not None
 
             assert entry.size == get_size_from_path(entry.instance_path)
+
+
+def test_set_store_state(
+    server, admin_client, librarian_database_session_maker, test_orm
+):
+    store_list = admin_client.get_store_list()
+
+    for store in store_list:
+        response = admin_client.set_store_state(store.name, enabled=False)
+
+        assert response is False
+
+        with librarian_database_session_maker() as session:
+            store = (
+                session.query(test_orm.StoreMetadata)
+                .filter_by(name=store.name)
+                .one_or_none()
+            )
+
+            assert not store.enabled
+
+        response = admin_client.set_store_state(store.name, enabled=True)
+
+        assert response is True
+
+        with librarian_database_session_maker() as session:
+            store = (
+                session.query(test_orm.StoreMetadata)
+                .filter_by(name=store.name)
+                .one_or_none()
+            )
+
+            assert store.enabled