Add devices to group tree

3 files changed, 103 insertions, 39 deletions

diff --git a/src/org/traccar/database/GroupTree.java b/src/org/traccar/database/GroupTree.java
index 7263a6331..b383b1501 100644
--- a/src/org/traccar/database/GroupTree.java
+++ b/src/org/traccar/database/GroupTree.java
@@ -15,6 +15,7 @@
  */
 package org.traccar.database;
 
+import org.traccar.model.Device;
 import org.traccar.model.Group;
 
 import java.util.ArrayList;
@@ -26,80 +27,126 @@ import java.util.Set;
 
 public class GroupTree {
 
-    private static class GroupNode {
+    private static class TreeNode {
 
         private Group group;
-        private GroupNode parent;
-        private Collection<GroupNode> children = new HashSet<>();
+        private Device device;
+        private Collection<TreeNode> children = new HashSet<>();
 
-        public GroupNode(Group group) {
+        public TreeNode(Group group) {
             this.group = group;
         }
 
+        public TreeNode(Device device) {
+            this.device = device;
+        }
+
         @Override
         public int hashCode() {
-            return (int) group.getId();
+            if (group != null) {
+                return (int) group.getId();
+            } else {
+                return (int) device.getId();
+            }
         }
 
         @Override
         public boolean equals(Object obj) {
-            if (!(obj instanceof GroupNode)) {
+            if (!(obj instanceof TreeNode)) {
                 return false;
             }
-            return obj == this || group.getId() == ((GroupNode) obj).group.getId();
+            TreeNode other = (TreeNode) obj;
+            if (other == this) {
+                return true;
+            }
+            if (group != null) {
+                if (other.group != null) {
+                    return group.getId() == other.group.getId();
+                }
+            } else if (device != null) {
+                if (other.device != null) {
+                    return device.getId() == other.device.getId();
+                }
+            }
+            return false;
         }
 
         public Group getGroup() {
             return group;
         }
 
-        public void setParent(GroupNode parent) {
-            this.parent = parent;
+        public Device getDevice() {
+            return device;
+        }
+
+        public void setParent(TreeNode parent) {
             if (parent != null) {
                 parent.children.add(this);
             }
         }
 
-        public GroupNode getParent() {
-            return parent;
-        }
-
-        public Collection<GroupNode> getChildren() {
+        public Collection<TreeNode> getChildren() {
             return children;
         }
 
     }
 
-    private final Map<Long, GroupNode> groupMap = new HashMap<>();
+    private final Map<Long, TreeNode> groupMap = new HashMap<>();
 
-    public GroupTree(Collection<Group> groups) {
+    public GroupTree(Collection<Group> groups, Collection<Device> devices) {
 
         for (Group group : groups) {
-            groupMap.put(group.getId(), new GroupNode(group));
+            groupMap.put(group.getId(), new TreeNode(group));
         }
 
-        for (GroupNode node : groupMap.values()) {
+        for (TreeNode node : groupMap.values()) {
             if (node.getGroup().getGroupId() != 0) {
                 node.setParent(groupMap.get(node.getGroup().getGroupId()));
             }
         }
 
+        Map<Long, TreeNode> deviceMap = new HashMap<>();
+
+        for (Device device : devices) {
+            deviceMap.put(device.getId(), new TreeNode(device));
+        }
+
+        for (TreeNode node : deviceMap.values()) {
+            if (node.getDevice().getGroupId() != 0) {
+                node.setParent(groupMap.get(node.getDevice().getGroupId()));
+            }
+        }
+
     }
 
-    public Collection<Group> getDescendants(long groupId) {
-        Set<GroupNode> results = new HashSet<>();
-        getDescendants(results, groupMap.get(groupId));
+    public Collection<Group> getGroups(long groupId) {
+        Set<TreeNode> results = new HashSet<>();
+        getNodes(results, groupMap.get(groupId));
         Collection<Group> groups = new ArrayList<>();
-        for (GroupNode node : results) {
-            groups.add(node.getGroup());
+        for (TreeNode node : results) {
+            if (node.getGroup() != null) {
+                groups.add(node.getGroup());
+            }
         }
         return groups;
     }
 
-    private void getDescendants(Set<GroupNode> results, GroupNode node) {
-        for (GroupNode child : node.getChildren()) {
+    public Collection<Device> getDevices(long groupId) {
+        Set<TreeNode> results = new HashSet<>();
+        getNodes(results, groupMap.get(groupId));
+        Collection<Device> devices = new ArrayList<>();
+        for (TreeNode node : results) {
+            if (node.getDevice() != null) {
+                devices.add(node.getDevice());
+            }
+        }
+        return devices;
+    }
+
+    private void getNodes(Set<TreeNode> results, TreeNode node) {
+        for (TreeNode child : node.getChildren()) {
             results.add(child);
-            getDescendants(results, child);
+            getNodes(results, child);
         }
     }
 
diff --git a/src/org/traccar/database/PermissionsManager.java b/src/org/traccar/database/PermissionsManager.java
index 0aefb9605..dc37bbc84 100644
--- a/src/org/traccar/database/PermissionsManager.java
+++ b/src/org/traccar/database/PermissionsManager.java
@@ -16,7 +16,6 @@
 package org.traccar.database;
 
 import java.sql.SQLException;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -69,23 +68,23 @@ public class PermissionsManager {
                 users.put(user.getId(), user);
             }
 
-            GroupTree groupTree = new GroupTree(dataManager.getAllGroups());
+            GroupTree groupTree = new GroupTree(dataManager.getAllGroups(), dataManager.getAllDevices());
             for (GroupPermission permission : dataManager.getGroupPermissions()) {
                 Set<Long> userGroupPermissions = getGroupPermissions(permission.getUserId());
+                Set<Long> userDevicePermissions = getDevicePermissions(permission.getUserId());
                 userGroupPermissions.add(permission.getGroupId());
-                for (Group group : groupTree.getDescendants(permission.getGroupId())) {
+                for (Group group : groupTree.getGroups(permission.getGroupId())) {
                     userGroupPermissions.add(group.getId());
                 }
+                for (Device device : groupTree.getDevices(permission.getGroupId())) {
+                    userDevicePermissions.add(device.getId());
+                }
             }
 
             for (DevicePermission permission : dataManager.getDevicePermissions()) {
                 getDevicePermissions(permission.getUserId()).add(permission.getDeviceId());
             }
 
-            for (Device device : dataManager.getAllDevices()) {
-                // TODO
-            }
-
         } catch (SQLException error) {
             Log.warning(error);
         }
diff --git a/test/org/traccar/database/GroupTreeTest.java b/test/org/traccar/database/GroupTreeTest.java
index 26ac595f8..f73cb75b6 100644
--- a/test/org/traccar/database/GroupTreeTest.java
+++ b/test/org/traccar/database/GroupTreeTest.java
@@ -2,6 +2,7 @@ package org.traccar.database;
 
 import org.junit.Assert;
 import org.junit.Test;
+import org.traccar.model.Device;
 import org.traccar.model.Group;
 
 import java.util.ArrayList;
@@ -16,7 +17,15 @@ public class GroupTreeTest {
         group.setGroupId(parent);
         return group;
     }
-    
+
+    private static Device createDevice(long id, String name, long parent) {
+        Device device = new Device();
+        device.setId(id);
+        device.setName(name);
+        device.setGroupId(parent);
+        return device;
+    }
+
     @Test
     public void testGetDescendants() {
         Collection<Group> groups = new ArrayList<>();
@@ -26,12 +35,21 @@ public class GroupTreeTest {
         groups.add(createGroup(4, "Fourth", 2));
         groups.add(createGroup(5, "Fifth", 4));
 
-        GroupTree groupTree = new GroupTree(groups);
+        Collection<Device> devices = new ArrayList<>();
+        devices.add(createDevice(1, "One", 3));
+        devices.add(createDevice(2, "Two", 5));
+        devices.add(createDevice(3, "One", 5));
+
+        GroupTree groupTree = new GroupTree(groups, devices);
+
+        Assert.assertEquals(4, groupTree.getGroups(1).size());
+        Assert.assertEquals(3, groupTree.getGroups(2).size());
+        Assert.assertEquals(0, groupTree.getGroups(3).size());
+        Assert.assertEquals(1, groupTree.getGroups(4).size());
 
-        Assert.assertEquals(4, groupTree.getDescendants(1).size());
-        Assert.assertEquals(3, groupTree.getDescendants(2).size());
-        Assert.assertEquals(0, groupTree.getDescendants(3).size());
-        Assert.assertEquals(1, groupTree.getDescendants(4).size());
+        Assert.assertEquals(3, groupTree.getDevices(1).size());
+        Assert.assertEquals(1, groupTree.getDevices(3).size());
+        Assert.assertEquals(2, groupTree.getDevices(4).size());
     }
 
 }