Added app data unit test and fixed downloading of app data files

Author: AriAlavi

main.py

@@ -1,6 +1,5 @@
 from __future__ import unicode_literals
 
-import pyperclip
 import asyncio
 import re
 import os
@@ -10,7 +9,7 @@
 from shutil import rmtree
 import traceback
 
-from rimlink import generateStructure, compareStructures, getRimworldConfigArea, isAdmin
+from rimlink import generateStructure, compareStructures, AppDataStructure, isAdmin
 
 
 
@@ -29,6 +28,7 @@ def validateIP(givenIp):
 
 
 def menu(prompt, validator, useClipboard=False):
+    import pyperclip
     if useClipboard:
         input(prompt)
         result = pyperclip.paste()
@@ -58,9 +58,10 @@ def hangForever():
     while True:
         time.sleep(120)
 
-def clientSyncFiles(to_delete, to_add, to_modify):
+def clientSyncFiles(to_delete, to_add, to_modify, **kwargs):
     to_delete.extend(to_modify)
     to_add.extend(to_modify)
+    testing = kwargs.get("testing", None)
     del to_modify
     folders = []
     for delete in to_delete:
@@ -86,6 +87,9 @@ def clientSyncFiles(to_delete, to_add, to_modify):
 
     i = 0
     for file_name in to_add:
+        if testing:
+            print("Downloaded {}".format(file_name))
+            continue
         s = socket.socket()
         s.connect((IP_ADDRESS, PORT))
         Server.clientRecieveFile(s, file_name.path())
@@ -139,7 +143,7 @@ def client():
     s = socket.socket()
     s.connect((IP_ADDRESS, PORT))
     if sync_config:
-        my_config = generateStructure(getRimworldConfigArea(), app_data=getRimworldConfigArea())
+        my_config = generateStructure(AppDataStructure.getRimworldConfigArea(), app_data=AppDataStructure.getRimworldConfigArea())
         my_config_pickled = pickle.dumps(my_config)
         s.send(b"\02")
         Server.clientSendPickle(s, my_config_pickled)
@@ -281,7 +285,7 @@ async def _handle_client(self, r, w):
     async def run(self):
         print("Analyzing rimworld...")
         self.base_structure = generateStructure(".")
-        self.base_app_data_structure = generateStructure(getRimworldConfigArea(), app_data=getRimworldConfigArea())
+        self.base_app_data_structure = generateStructure(AppDataStructure.getRimworldConfigArea(), app_data=AppDataStructure.getRimworldConfigArea())
         print("Ready to receive connections on {}:{}".format(IP_ADDRESS, PORT))
         await asyncio.start_server(self._handle_client, IP_ADDRESS, PORT)
 

rimlink.py

@@ -15,11 +15,6 @@ def isAdmin():
         is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0
     return is_admin
 
-def getRimworldConfigArea():
-    roaming = os.getenv("APPDATA")
-    app_data = "\\".join(roaming.split("\\")[:-1])
-    return os.path.join(os.path.join(os.path.join(app_data, "LocalLow"), "Ludeon Studios"), "RimWorld by Ludeon Studios")
-
 class FileFolder:
     def __init__(self, name, parent=None, **kwargs):
         assert isinstance(name, str)
@@ -64,12 +59,17 @@ def hashFile(givenFile):
     h  = hashlib.sha256()
     b  = bytearray(128*1024)
     mv = memoryview(b)
-    try:
-        with open(givenFile, 'rb', buffering=0) as f:
-            for n in iter(lambda : f.readinto(mv), 0):
-                h.update(mv[:n])
-    except:
-        return "permission_denied"
+    if os.path.isdir(givenFile):
+        return "folder"
+    elif os.path.isfile(givenFile):
+        try:
+            with open(givenFile, 'rb', buffering=0) as f:
+                for n in iter(lambda : f.readinto(mv), 0):
+                    h.update(mv[:n])
+        except:
+            return "permission_denied"
+    else:
+        raise Exception("{} does not exist".format(givenFile))
     return h.hexdigest()
 
 def compareFiles(file1, file2):
@@ -86,14 +86,19 @@ def __init__(self, name, parent=None, **kwargs):
 
 class AppDataStructure(HashStructure):
     def __init__(self, name, parent=None, **kwargs):
-        self.app_data_path = kwargs['app_data']
         super(AppDataStructure, self).__init__(name, parent, **kwargs)
 
     def path(self):
         if self.parent:
             return os.path.join(self.parent.path(), self.name)
         else:
-            return self.app_data_path
+            return AppDataStructure.getRimworldConfigArea()
+    @staticmethod
+    def getRimworldConfigArea():
+        roaming = os.getenv("APPDATA")
+        app_data = "\\".join(roaming.split("\\")[:-1])
+        return os.path.join(os.path.join(os.path.join(app_data, "LocalLow"), "Ludeon Studios"), "RimWorld by Ludeon Studios")
+
 
 FILE_EXCEPTIONS = ["__pycache__", "Saves", "Scenarios", "MpReplays", "MpDesyncs", "Player.log", "Player-prev.log", ".gitignore", ".git", "rimlink.exe"]
 def generateStructure(relativePositionStart, parent=None, **kwargs):

test.py

@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 
 from rimlink import *
+from main import clientSyncFiles
 
 import unittest
 import os
@@ -142,7 +143,7 @@ def test_different_deep(self):
         self.assertEqual(results['modify'], [r"Interior\deep\hihi.txt"])
         self.assertEqual(results['add'], [])  
     def test_app_data(self):
-        TOP_LOCATION = getRimworldConfigArea()
+        TOP_LOCATION = AppDataStructure.getRimworldConfigArea()
         INNER_LOCATION = os.path.join(TOP_LOCATION, "testing")
         try:
             os.mkdir(INNER_LOCATION)
@@ -156,12 +157,12 @@ def test_app_data(self):
         file = open(add_inner, "w")
         file.write("testing")
         file.close()
-        base_structure = generateStructure(getRimworldConfigArea(), app_data=getRimworldConfigArea())
+        base_structure = generateStructure(AppDataStructure.getRimworldConfigArea(), app_data=AppDataStructure.getRimworldConfigArea())
 
         os.remove(add)
         os.remove(add_inner)
 
-        after_structure = generateStructure(getRimworldConfigArea(), app_data=getRimworldConfigArea())
+        after_structure = generateStructure(AppDataStructure.getRimworldConfigArea(), app_data=AppDataStructure.getRimworldConfigArea())
 
         results = compareStructures(base_structure, after_structure)
 
@@ -177,9 +178,22 @@ def test_app_data(self):
     def test_different_app_data(self):
         APP_DATA_BASE = "test_files/FakeAppData1/"
         APP_DATA_DIFFERENT = "test_files/FakeAppData2/"
+
+        file_name = "different.txt"
+        file = open(os.path.join(APP_DATA_BASE, file_name), "w")
+        file.write("good")
+        file.close()
+        file = open(os.path.join(APP_DATA_DIFFERENT, file_name), "w")
+        file.write("bad")
+        file.close()
+        AppDataStructure.getRimworldConfigArea = lambda : APP_DATA_BASE
         base_structure = generateStructure(APP_DATA_BASE, app_data=APP_DATA_BASE)
+        AppDataStructure.getRimworldConfigArea = lambda : APP_DATA_DIFFERENT
         different_structure = generateStructure(APP_DATA_DIFFERENT, app_data=APP_DATA_DIFFERENT)
+        
         results = compareStructures(base_structure, different_structure)
+        self.assertEqual(results['modify'][0].path(), 'test_files/FakeAppData2/different.txt')
+        clientSyncFiles(results['delete'], results['add'], results['modify'], testing=True)
 
 
 class IsFileTest(unittest.TestCase):
@@ -196,7 +210,7 @@ def test_is_not_file(self):
         self.assertFalse(FileFolder("test_files\\RimworldBase\\empty").file)
 
     def test_file_indirect(self):
-        config_location = getRimworldConfigArea()
+        config_location = AppDataStructure.getRimworldConfigArea()
         parent = generateStructure(config_location, app_data=config_location)
         self.assertFalse(parent.file)
         for child in parent.children: