From ebabea39ce885e2f88584337d51d55e8153844f9 Mon Sep 17 00:00:00 2001
From: MMGen <mmgen@tuta.io>
Date: Mon, 26 Feb 2018 12:34:59 +0000
Subject: [PATCH] tw.add_label(): code cleanups, tests

---
 mmgen/rpc.py  |  3 +++
 mmgen/tool.py | 10 +++++-----
 mmgen/tw.py   | 27 +++++++++------------------
 test/test.py  | 45 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 23 deletions(-)

diff --git a/mmgen/rpc.py b/mmgen/rpc.py
index 02a2d1f5..5db1381b 100755
--- a/mmgen/rpc.py
+++ b/mmgen/rpc.py
@@ -71,6 +71,9 @@ class CoinDaemonRPCConnection(object):
 	# With on_fail='return', returns 'rpcfail',(resp_object,(die_args))
 	def request(self,cmd,*args,**kwargs):
 
+		if os.getenv('MMGEN_RPC_FAIL_ON_COMMAND') == cmd:
+			cmd = 'badcommand_' + cmd
+
 		cf = { 'timeout':g.http_timeout, 'batch':False, 'on_fail':'die' }
 
 		for k in cf:
diff --git a/mmgen/tool.py b/mmgen/tool.py
index aa67d8e8..2b7a054e 100755
--- a/mmgen/tool.py
+++ b/mmgen/tool.py
@@ -83,8 +83,8 @@ cmd_data = OrderedDict([
 	('Txview',       ['<{} TX file(s)> [str]'.format(pnm),'pager [bool=False]','terse [bool=False]',"sort [str='mtime'] (options: ctime, atime)",'MARGS']),
 	('Twview',       ["sort [str='age']",'reverse [bool=False]','show_days [bool=True]','show_mmid [bool=True]','minconf [int=1]','wide [bool=False]','pager [bool=False]']),
 
-	('Add_label',       ['<{} address> [str]'.format(pnm),'<label> [str]']),
-	('Remove_label',    ['<{} address> [str]'.format(pnm)]),
+	('Add_label',       ['<{} or coin address> [str]'.format(pnm),'<label> [str]']),
+	('Remove_label',    ['<{} or coin address> [str]'.format(pnm)]),
 	('Addrfile_chksum', ['<{} addr file> [str]'.format(pnm),"mmtype [str='']"]),
 	('Keyaddrfile_chksum', ['<{} addr file> [str]'.format(pnm),"mmtype [str='']"]),
 	('Passwdfile_chksum', ['<{} password file> [str]'.format(pnm)]),
@@ -826,9 +826,9 @@ def Twview(pager=False,reverse=False,wide=False,minconf=1,sort='age',show_days=T
 	out = tw.format_for_printing(color=True) if wide else tw.format_for_display()
 	(Msg_r,do_pager)[pager](out)
 
-def Add_label(mmaddr,label):
+def Add_label(mmaddr_or_coin_addr,label):
 	rpc_init()
 	from mmgen.tw import MMGenTrackingWallet
-	MMGenTrackingWallet.add_label(mmaddr,label) # dies on failure
+	MMGenTrackingWallet.add_label(mmaddr_or_coin_addr,label) # dies on failure
 
-def Remove_label(mmaddr): Add_label(mmaddr,'')
+def Remove_label(mmaddr_or_coin_addr): Add_label(mmaddr_or_coin_addr,'')
diff --git a/mmgen/tw.py b/mmgen/tw.py
index c056a7d4..81002211 100755
--- a/mmgen/tw.py
+++ b/mmgen/tw.py
@@ -321,33 +321,24 @@ Display options: show [D]ays, [g]roup, show [m]mgen addr, r[e]draw screen
 		if is_mmgen_id(arg1):
 			mmaddr = TwMMGenID(arg1)
 
-		if not coinaddr and not mmaddr:
-			msg("Address '{}' invalid or not found in tracking wallet".format(addr or arg1))
-			return False
-
-		if not coinaddr:
+		if mmaddr and not coinaddr:
 			from mmgen.addr import AddrData
 			coinaddr = AddrData(source='tw').mmaddr2coinaddr(mmaddr)
 
-		if not coinaddr:
-			msg("{} address '{}' not found in tracking wallet".format(g.proj_name,mmaddr))
-			return False
-
-		# Checked that the user isn't importing a random address
-		if not coinaddr.is_in_tracking_wallet():
-			msg("Address '{}' not in tracking wallet".format(coinaddr))
-			return False
-
-		if not coinaddr.is_for_chain(g.chain):
-			msg("Address '{}' not valid for chain {}".format(coinaddr,g.chain.upper()))
+		try:
+			if not is_mmgen_id(arg1):
+				assert coinaddr,"Invalid coin address for this chain: {}".format(arg1)
+			assert coinaddr,"{pn} address '{ma}' not found in tracking wallet"
+			assert coinaddr.is_in_tracking_wallet(),"Address '{ca}' not found in tracking wallet"
+		except Exception as e:
+			msg(e[0].format(pn=g.proj_name,ma=mmaddr,ca=coinaddr))
 			return False
 
 		# Allow for the possibility that BTC addr of MMGen addr was entered.
 		# Do reverse lookup, so that MMGen addr will not be marked as non-MMGen.
 		if not mmaddr:
 			from mmgen.addr import AddrData
-			ad = AddrData(source='tw')
-			mmaddr = ad.coinaddr2mmaddr(coinaddr)
+			mmaddr = AddrData(source='tw').coinaddr2mmaddr(coinaddr)
 
 		if not mmaddr: mmaddr = '{}:{}'.format(g.proto.base_coin.lower(),coinaddr)
 
diff --git a/test/test.py b/test/test.py
index 0168a87e..b29cc85c 100755
--- a/test/test.py
+++ b/test/test.py
@@ -687,6 +687,13 @@ cmd_group['regtest'] = (
 	('regtest_alice_remove_label1','removing a label'),
 	('regtest_alice_chk_label4',   'the label'),
 	('regtest_alice_send_estimatefee','tx creation with no fee on command line'),
+	('regtest_alice_add_label_coinaddr','adding a label using the coin address'),
+	('regtest_alice_chk_label_coinaddr','the label'),
+	('regtest_alice_add_label_badaddr1','adding a label with invalid address'),
+	('regtest_alice_add_label_badaddr2','adding a label with invalid address for this chain'),
+	('regtest_alice_add_label_badaddr3','adding a label with wrong MMGen address'),
+	('regtest_alice_add_label_badaddr4','adding a label with wrong coin address'),
+	('regtest_alice_add_label_rpcfail','RPC failure code'),
 	('regtest_stop',               'stopping regtest daemon'),
 )
 
@@ -2506,6 +2513,44 @@ class MMGenTestSuite(object):
 		sid = self.regtest_user_sid('alice')
 		return self.regtest_user_add_label(name,'alice',sid+':C:1','Replacement Label')
 
+	def regtest_alice_add_label_coinaddr(self,name):
+		mmaddr = self.regtest_user_sid('alice') + ':C:2'
+		t = MMGenExpect(name,'mmgen-tool',['--alice','listaddress',mmaddr],no_msg=True)
+		btcaddr = [i for i in t.read().splitlines() if i.lstrip()[0:len(mmaddr)] == mmaddr][0].split()[1]
+		return self.regtest_user_add_label(name,'alice',btcaddr,'Label added using coin address')
+
+	def regtest_alice_chk_label_coinaddr(self,name):
+		sid = self.regtest_user_sid('alice')
+		return self.regtest_user_chk_label(name,'alice',sid+':C:2','Label added using coin address')
+
+	def regtest_alice_add_label_badaddr(self,name,addr,reply):
+		t = MMGenExpect(name,'mmgen-tool',['--alice','add_label',addr,'(none)'])
+		t.expect(reply,regex=True)
+		t.ok()
+
+	def regtest_alice_add_label_badaddr1(self,name):
+		return self.regtest_alice_add_label_badaddr(name,'abc','Invalid coin address for this chain: abc')
+
+	def regtest_alice_add_label_badaddr2(self,name):
+		addr = g.proto.pubhash2addr('00'*20,False) # mainnet zero address
+		return self.regtest_alice_add_label_badaddr(name,addr,'Invalid coin address for this chain: '+addr)
+
+	def regtest_alice_add_label_badaddr3(self,name):
+		addr = self.regtest_user_sid('alice') + ':C:123'
+		return self.regtest_alice_add_label_badaddr(name,addr,
+			"MMGen address '{}' not found in tracking wallet".format(addr))
+
+	def regtest_alice_add_label_badaddr4(self,name):
+		addr = CoinProtocol(g.coin,True).pubhash2addr('00'*20,False) # testnet zero address
+		return self.regtest_alice_add_label_badaddr(name,addr,
+			"Address '{}' not found in tracking wallet".format(addr))
+
+	def regtest_alice_add_label_rpcfail(self,name):
+		addr = self.regtest_user_sid('alice') + ':C:2'
+		os.environ['MMGEN_RPC_FAIL_ON_COMMAND'] = 'importaddress'
+		self.regtest_alice_add_label_badaddr(name,addr,'Label could not be added')
+		os.environ['MMGEN_RPC_FAIL_ON_COMMAND'] = ''
+
 	def regtest_alice_remove_label1(self,name):
 		sid = self.regtest_user_sid('alice')
 		return self.regtest_user_remove_label(name,'alice',sid+':C:1')