From e36c0a6becb99888878897cd40a6d6427a30dda7 Mon Sep 17 00:00:00 2001
From: arkohut <39525455+arkohut@users.noreply.github.com>
Date: Tue, 4 Jun 2024 17:03:31 +0800
Subject: [PATCH] feat(folder): prevent duplicated folders

---
 memos/server.py      |  7 +++++++
 memos/test_server.py | 40 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/memos/server.py b/memos/server.py
index b51f753..e9f8928 100644
--- a/memos/server.py
+++ b/memos/server.py
@@ -41,6 +41,9 @@ def root():
 
 @app.post("/libraries", response_model=Library)
 def new_library(library_param: NewLibraryParam, db: Session = Depends(get_db)):
+    # Remove duplicate folders from the library_param
+    unique_folders = list(set(library_param.folders))
+    library_param.folders = unique_folders
     library = crud.create_library(library_param, db)
     return library
 
@@ -68,6 +71,10 @@ def new_folder(
     library = crud.get_library_by_id(library_id, db)
     if library is None:
         raise HTTPException(status_code=404, detail="Library not found")
+    
+    existing_folders = [folder.path for folder in library.folders]
+    if str(folder.path) in existing_folders:
+        raise HTTPException(status_code=400, detail="Folder already exists in the library")
 
     return crud.add_folder(library_id=library.id, folder=folder, db=db)
 
diff --git a/memos/test_server.py b/memos/test_server.py
index b87a25c..4a2bf4e 100644
--- a/memos/test_server.py
+++ b/memos/test_server.py
@@ -260,7 +260,7 @@ def test_get_entity_by_filepath(client):
 def test_list_entities_in_folder(client):
     # Setup data: Create a new library and folder
     new_library = NewLibraryParam(
-        name="Library for List Entities Test", folders=["/tmp"]
+        name="Library for List Entities Test", folders=[]
     )
     library_response = client.post(
         "/libraries", json=new_library.model_dump(mode="json")
@@ -355,3 +355,41 @@ def test_remove_entity(client):
     invalid_delete_response = client.delete(f"/libraries/{library_id}/entities/9999")
     assert invalid_delete_response.status_code == 404
     assert invalid_delete_response.json() == {"detail": "Entity not found in the specified library"}
+
+
+def test_add_folder_to_library(client):
+
+    import os
+
+    # Prepare tmp folders for the test
+    tmp_folder_path = "/tmp/new_folder"
+    if not os.path.exists(tmp_folder_path):
+        os.makedirs(tmp_folder_path)
+
+    # Create a new library
+    new_library = NewLibraryParam(name="Test Library")
+    library_response = client.post("/libraries", json=new_library.model_dump(mode="json"))
+    library_id = library_response.json()["id"]
+
+    # Add a new folder to the library
+    new_folder = NewFolderParam(path="/tmp/new_folder")
+    folder_response = client.post(f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json"))
+    assert folder_response.status_code == 200
+    assert folder_response.json()["path"] == "/tmp/new_folder"
+
+    # Verify the folder is added
+    library_response = client.get(f"/libraries/{library_id}")
+    assert library_response.status_code == 200
+    library_data = library_response.json()
+    folder_paths = [folder["path"] for folder in library_data["folders"]]
+    assert "/tmp/new_folder" in folder_paths
+
+    # Test for adding a folder that already exists
+    duplicate_folder_response = client.post(f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json"))
+    assert duplicate_folder_response.status_code == 400
+    assert duplicate_folder_response.json() == {"detail": "Folder already exists in the library"}
+
+    # Test for adding a folder to a non-existent library
+    invalid_folder_response = client.post(f"/libraries/9999/folders", json=new_folder.model_dump(mode="json"))
+    assert invalid_folder_response.status_code == 404
+    assert invalid_folder_response.json() == {"detail": "Library not found"}