final changes

PhoneBook.java

- 
+//JN 2018  
 import java.util.*;
-/**
- * Created by leon on 1/23/18.
- */
 public class PhoneBook {
     private TreeMap<String, ArrayList<String>> phoneBook;
     public PhoneBook() {
     public void add(String name, String phoneNumber) {
         ArrayList<String> input = phoneBook.containsKey(name) ? 
                                   phoneBook.get(name)
-                                  :                         
-                                  new ArrayList<String>();
+                                : new ArrayList<String>();
         input.add(phoneNumber);
         phoneBook.put(name, input);
     }
     
+    public void remove (String phoneNumber) {
+        ArrayList<String> edit = new ArrayList<String>();
+        Set<String> keys = phoneBook.keySet();
+        for (String key : keys) {
+            if (phoneBook.get(key).contains(phoneNumber)) {
+                edit = phoneBook.get(key);
+                edit.remove(edit.indexOf(phoneNumber));
+                phoneBook.put(key, edit);
+            }
+        }
+    }
+    
     public void removeEntry(String name) {
         phoneBook.remove(name);
     }
     
     public String lookup(String name) {
-        return phoneBook.get(name)
+        
+        return phoneBook.containsKey(name) ?
+               phoneBook.get(name)
                         .toString()
-                        .replaceAll("[\\[\\]]", "");
+                        .replaceAll("[\\[\\]]", "")
+            : "N/A";
     }
     
     public String reverseLookup(String phoneNumber) {
                 output = key;
             }
         }
-        return output.equals("") ? "No associated name." : output;
+        return output.equals("") ? "N/A" : output;
     }
     
     public String display() {

PhoneBookTest.java

        assertEquals(expectedNumber, actual);
     }
     
-    @Test(expected = NullPointerException.class)
+    @Test
     public void testRemoveEntry_And_Lookup()
     {
        PhoneBook test = new PhoneBook();
        String number = "(302)-867-5309";
        test.add(name, number);
        test.removeEntry(name);
-       test.lookup(name);   
+       String actual = test.lookup(name);
+       assertEquals("N/A", actual);
+    }
+    
+    @Test
+    public void testRemoveIndividualNumber() {
+       PhoneBook test = new PhoneBook();
+       String name = "Jenny";
+       String number = "(302)-867-5309";
+       String number2 = "(601)-482-9832";
+       test.add(name, number);
+       test.add(name, number2);
+       test.remove(number);
+       String actual = test.reverseLookup(number);
+       assertEquals("N/A", actual);
     }
     
     @Test
        test.add("hey", "302");
        test.add("yo", "521");
        test.add("hey", "498");
-       System.out.println(test.display());
+       test.display();
     }
 }