Store owner information on creating a note and added integration test

cases for the relevant scenarios.
apache · Sep 1, 2016 · d794a2a · d794a2a
1 parent 323aa18
commit d794a2a
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 2 deletions.
diff --git a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/Notebook.java b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/Notebook.java
@@ -24,10 +24,12 @@
 import java.util.Comparator;
 import java.util.Date;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 import com.google.gson.Gson;
@@ -158,6 +160,11 @@ public Note createNote(List<String> interpreterIds, AuthenticationInfo subject)
       note.putDefaultReplName();
     }
 
+    if (subject != null && !"anonymous".equals(subject.getUser())) {
+      Set<String> owners = new HashSet<String>();
+      owners.add(subject.getUser());
+      notebookAuthorization.setOwners(note.getId(), owners);
+    }
     notebookIndex.addIndexDoc(note);
     note.persist(subject);
     fireNoteCreateEvent(note);

diff --git a/zeppelin-zengine/src/test/java/org/apache/zeppelin/notebook/NotebookTest.java b/zeppelin-zengine/src/test/java/org/apache/zeppelin/notebook/NotebookTest.java
@@ -42,6 +42,7 @@
 import org.apache.zeppelin.scheduler.Job.Status;
 import org.apache.zeppelin.scheduler.SchedulerFactory;
 import org.apache.zeppelin.search.SearchService;
+import org.apache.zeppelin.user.AuthenticationInfo;
 import org.apache.zeppelin.user.Credentials;
 import org.junit.After;
 import org.junit.Before;
@@ -209,6 +210,21 @@ public void testPersist() throws IOException, SchedulerException, RepositoryExce
     assertEquals(1, notebook2.getAllNotes().size());
   }
 
+  @Test
+  public void testCreateNoteWithSubject() throws IOException, SchedulerException, RepositoryException {
+    AuthenticationInfo subject = new AuthenticationInfo("user1");
+    Note note = notebook.createNote(subject);
+
+    Notebook notebook = new Notebook(
+        conf, notebookRepo, schedulerFactory,
+        new InterpreterFactory(conf, null, null, null, depResolver), this, null, notebookAuthorization, null);
+    assertNotNull(notebook.getNotebookAuthorization().getOwners(note.getId()));
+    assertEquals(1, notebook.getNotebookAuthorization().getOwners(note.getId()).size());
+    Set<String> owners = new HashSet<>();
+    owners.add("user1");
+    assertEquals(owners, notebook.getNotebookAuthorization().getOwners(note.getId()));
+  }
+
   @Test
   public void testClearParagraphOutput() throws IOException, SchedulerException{
     Note note = notebook.createNote(null);
@@ -351,7 +367,7 @@ public void testAutoRestartInterpreterAfterSchedule() throws InterruptedExceptio
 
   @Test
   public void testExportAndImportNote() throws IOException, CloneNotSupportedException,
-          InterruptedException {
+          InterruptedException, InterpreterException, SchedulerException, RepositoryException {
     Note note = notebook.createNote(null);
     factory.setInterpreters(note.getId(), factory.getDefaultInterpreterSettingList());
 
@@ -374,11 +390,23 @@ public void testExportAndImportNote() throws IOException, CloneNotSupportedExcep
     assertEquals(p.getId(), p2.getId());
     assertEquals(p.text, p2.text);
     assertEquals(p.getResult().message(), p2.getResult().message());
+
+    // Verify import note with subject
+    AuthenticationInfo subject = new AuthenticationInfo("user1");
+    Note importedNote2 = notebook.importNote(exportedNoteJson, "Title2", subject);
+    Notebook notebook = new Notebook(
+        conf, notebookRepo, schedulerFactory,
+        new InterpreterFactory(conf, null, null, null, depResolver), this, null, notebookAuthorization, null);
+    assertNotNull(notebook.getNotebookAuthorization().getOwners(importedNote2.getId()));
+    assertEquals(1, notebook.getNotebookAuthorization().getOwners(importedNote2.getId()).size());
+    Set<String> owners = new HashSet<>();
+    owners.add("user1");
+    assertEquals(owners, notebook.getNotebookAuthorization().getOwners(importedNote2.getId()));
   }
 
   @Test
   public void testCloneNote() throws IOException, CloneNotSupportedException,
-      InterruptedException {
+      InterruptedException, InterpreterException, SchedulerException, RepositoryException {
     Note note = notebook.createNote(null);
     factory.setInterpreters(note.getId(), factory.getDefaultInterpreterSettingList());
 
@@ -396,6 +424,18 @@ public void testCloneNote() throws IOException, CloneNotSupportedException,
     assertEquals(cp.getId(), p.getId());
     assertEquals(cp.text, p.text);
     assertEquals(cp.getResult().message(), p.getResult().message());
+
+    // Verify clone note with subject
+    AuthenticationInfo subject = new AuthenticationInfo("user1");
+    Note cloneNote2 = notebook.cloneNote(note.getId(), "clone note2", subject);
+    Notebook notebook = new Notebook(
+        conf, notebookRepo, schedulerFactory,
+        new InterpreterFactory(conf, null, null, null, depResolver), this, null, notebookAuthorization, null);
+    assertNotNull(notebook.getNotebookAuthorization().getOwners(cloneNote2.getId()));
+    assertEquals(1, notebook.getNotebookAuthorization().getOwners(cloneNote2.getId()).size());
+    Set<String> owners = new HashSet<>();
+    owners.add("user1");
+    assertEquals(owners, notebook.getNotebookAuthorization().getOwners(cloneNote2.getId()));
   }
 
   @Test