From 43699a74e43380c6e08a7cbc100576925a8a0291 Mon Sep 17 00:00:00 2001
From: Alexander Song <axiomofjoy@gmail.com>
Date: Thu, 12 Sep 2024 14:53:29 -0700
Subject: [PATCH] add it

---
 integration_tests/_helpers.py       | 11 +++++++++++
 integration_tests/auth/test_auth.py | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/integration_tests/_helpers.py b/integration_tests/_helpers.py
index f0dd4bb3c5..fabdd3eaee 100644
--- a/integration_tests/_helpers.py
+++ b/integration_tests/_helpers.py
@@ -211,6 +211,9 @@ def create_api_key(
     def delete_api_key(self, api_key: _ApiKey, /) -> None:
         return _delete_api_key(api_key, self)
 
+    def export_embeddings(self, filename: str) -> None:
+        _export_embeddings(self, filename=filename)
+
 
 _SYSTEM_USER_GID = _GqlId(GlobalID(type_name="User", node_id="1"))
 _DEFAULT_ADMIN = _User(
@@ -828,6 +831,13 @@ def _log_out(
     resp.raise_for_status()
 
 
+def _export_embeddings(
+    auth: Optional[_SecurityArtifact] = None, /, *, filename: str
+) -> httpx.Response:
+    resp = _httpx_client(auth).get("/exports", params={"filename": filename})
+    resp.raise_for_status()
+
+
 def _json(
     resp: httpx.Response,
 ) -> Dict[str, Any]:
@@ -853,3 +863,4 @@ def __exit__(self, *args: Any, **kwargs: Any) -> None: ...
 
 _EXPECTATION_401 = pytest.raises(HTTPStatusError, match="401 Unauthorized")
 _EXPECTATION_403 = pytest.raises(HTTPStatusError, match="403 Forbidden")
+_EXPECTATION_404 = pytest.raises(HTTPStatusError, match="404 Not Found")
diff --git a/integration_tests/auth/test_auth.py b/integration_tests/auth/test_auth.py
index 3d5e0deb91..4f5444fc1f 100644
--- a/integration_tests/auth/test_auth.py
+++ b/integration_tests/auth/test_auth.py
@@ -36,6 +36,7 @@
     _DEFAULT_ADMIN,
     _DENIED,
     _EXPECTATION_401,
+    _EXPECTATION_404,
     _MEMBER,
     _OK,
     _OK_OR_DENIED,
@@ -45,6 +46,7 @@
     _create_user,
     _DefaultAdminTokenSequestration,
     _Expectation,
+    _export_embeddings,
     _GetUser,
     _GqlId,
     _Headers,
@@ -741,3 +743,20 @@ def test_api_key(
         if api_key and expected is SpanExportResult.SUCCESS:
             _DEFAULT_ADMIN.delete_api_key(api_key)
             assert export(_spans) is SpanExportResult.FAILURE
+
+
+class TestEmbeddingsRestApi:
+    @pytest.mark.parametrize("role_or_user", [_MEMBER, _ADMIN, _DEFAULT_ADMIN])
+    def test_authenticated_users_can_access_route(
+        self,
+        role_or_user: _RoleOrUser,
+        _get_user: _GetUser,
+    ) -> None:
+        user = _get_user(role_or_user)
+        logged_in_user = user.log_in()
+        with _EXPECTATION_404:  # no files have been exported
+            logged_in_user.export_embeddings("embeddings")
+
+    def test_unauthenticated_requests_receive_401(self) -> None:
+        with _EXPECTATION_401:
+            _export_embeddings(None, filename="embeddings")