Full Coverage on Contracts

StorjOld · Jun 24, 2015 · 4d4dabb · 4d4dabb
1 parent b0b99f4
commit 4d4dabb
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 deletions.
diff --git a/dataserv/Contract.py b/dataserv/Contract.py
@@ -33,10 +33,18 @@ def new_contract(self, btc_addr, seed=None, byte_size=None):
         self.btc_addr = btc_addr
         self.contract_type = 0
 
-        # generate a random seed
-        seed = os.urandom(12)
-        self.seed = binascii.hexlify(seed).decode('ascii')
+        # take in a seed, if not generate it ourselves
+        if seed is None:
+            seed = os.urandom(12)
+            self.seed = binascii.hexlify(seed).decode('ascii')
+        else:
+            self.seed = seed
+
+        # take in a byte_size, if not then get it from config
+        if byte_size is None:
+            self.byte_size = app.config["SHARD_SIZE"]
+        else:
+            self.byte_size = byte_size
 
-        self.byte_size = app.config["SHARD_SIZE"]
         gen_file = RandomIO.RandomIO(seed).read(self.byte_size)
         self.file_hash = hashlib.sha256(gen_file).hexdigest()
diff --git a/tests/test_Contract.py b/tests/test_Contract.py
@@ -15,4 +15,28 @@ def test_new_contract(self):
         addr = '191GVvAaTRxLmz3rW3nU5jAV1rF186VxQc'
         my_contract = Contract()
         my_contract.new_contract(addr, 'ba17da75c580a0749b6c3d32', 1024)
-        print(my_contract.to_json())
+
+        self.assertEqual(my_contract.btc_addr, addr)
+        self.assertEqual(my_contract.contract_type, 0)
+        self.assertEqual(my_contract.file_hash, '8366600f680255341531972a68f201c9edee89350e6d8c43c2c22a385bf02a60')
+        self.assertEqual(my_contract.byte_size, 1024)
+        self.assertEqual(my_contract.seed, 'ba17da75c580a0749b6c3d32')
+
+    def test_new_contract2(self):
+        addr = '191GVvAaTRxLmz3rW3nU5jAV1rF186VxQc'
+        my_contract = Contract()
+        my_contract.new_contract(addr)
+        self.assertEqual(my_contract.btc_addr, addr)
+        self.assertEqual(my_contract.contract_type, 0)
+
+    def test_new_contract_json(self):
+        addr = '191GVvAaTRxLmz3rW3nU5jAV1rF186VxQc'
+        my_contract = Contract()
+        my_contract.new_contract(addr, 'ba17da75c580a0749b6c3d32', 1024)
+        json_contract = my_contract.to_json()
+
+        self.assertEqual(json_contract["btc_addr"], addr)
+        self.assertEqual(json_contract["contract_type"], 0)
+        self.assertEqual(json_contract["file_hash"], '8366600f680255341531972a68f201c9edee89350e6d8c43c2c22a385bf02a60')
+        self.assertEqual(json_contract["byte_size"], 1024)
+        self.assertEqual(json_contract["seed"], 'ba17da75c580a0749b6c3d32')