From bb76e1e3fa2983083348eb20c91c80b757226ba0 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Sat, 13 Jun 2020 17:42:11 +0000
Subject: [PATCH] addrimport: display token info; TrackingWallet: token import
 fix

---
 mmgen/altcoins/eth/tw.py    | 15 ++++++++-------
 mmgen/main_addrimport.py    |  4 ++++
 test/test_py_d/ts_ethdev.py | 37 +++++++++++++++++++++++--------------
 3 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/mmgen/altcoins/eth/tw.py b/mmgen/altcoins/eth/tw.py
index 8b21ee6d..f0ba4abb 100755
--- a/mmgen/altcoins/eth/tw.py
+++ b/mmgen/altcoins/eth/tw.py
@@ -170,13 +170,14 @@ class EthereumTokenTrackingWallet(EthereumTrackingWallet):
 		from mmgen.obj import TokenAddr
 		self.token = TokenAddr(proto,token_addr)
 
-		if self.token in self.data['tokens']:
-			self.decimals = self.get_param('decimals')
-			self.symbol   = self.get_param('symbol')
-		elif self.importing:
-			await self.import_token(self.token) # sets self.decimals, self.symbol
-		else:
-			raise TokenNotInWallet(f'Specified token {self.token!r} not in wallet!')
+		if self.token not in self.data['tokens']:
+			if self.importing:
+				await self.import_token(self.token)
+			else:
+				raise TokenNotInWallet(f'Specified token {self.token!r} not in wallet!')
+
+		self.decimals = self.get_param('decimals')
+		self.symbol   = self.get_param('symbol')
 
 		proto.tokensym = self.symbol
 
diff --git a/mmgen/main_addrimport.py b/mmgen/main_addrimport.py
index 808032a9..0be3a2f0 100755
--- a/mmgen/main_addrimport.py
+++ b/mmgen/main_addrimport.py
@@ -158,11 +158,15 @@ async def main():
 	from .tw import TrackingWallet
 	if opt.token_addr:
 		proto.tokensym = 'foo' # hack to trigger 'Token' in altcoin_subclass()
+
 	tw = await TrackingWallet(
 		proto      = proto,
 		token_addr = opt.token_addr,
 		mode       = 'i' )
 
+	if opt.token or opt.token_addr:
+		msg(f'Importing for token {tw.token.hl()} ({tw.token.hlc(proto.tokensym)})')
+
 	from .rpc import rpc_init
 	tw.rpc = await rpc_init(proto)
 
diff --git a/test/test_py_d/ts_ethdev.py b/test/test_py_d/ts_ethdev.py
index 5537c08d..c75335c4 100755
--- a/test/test_py_d/ts_ethdev.py
+++ b/test/test_py_d/ts_ethdev.py
@@ -216,8 +216,10 @@ class TestSuiteEthdev(TestSuiteBase,TestSuiteShared):
 		('token_addrgen',        'generating token addresses'),
 		('token_addrimport_badaddr1','importing token addresses (no token address)'),
 		('token_addrimport_badaddr2','importing token addresses (bad token address)'),
-		('token_addrimport',     'importing token addresses'),
-		('token_addrimport_batch','importing token addresses (dummy batch mode)'),
+		('token_addrimport_addr1','importing token addresses using token address (MM1)'),
+		('token_addrimport_addr2','importing token addresses using token address (MM2)'),
+		('token_addrimport_batch','importing token addresses (dummy batch mode) (MM1)'),
+		('token_addrimport_sym',  'importing token addresses using token symbol (MM2)'),
 
 		('bal7',                'the {} balance'.format(coin)),
 		('token_bal1',          'the {} balance and token balance'.format(coin)),
@@ -729,20 +731,27 @@ class TestSuiteEthdev(TestSuiteBase,TestSuiteShared):
 		t.req_exit_val = 2
 		return t
 
-	def token_addrimport(self,extra_args=[],expect='3/3'):
-		for n,r in ('1','11-13'),('2','21-23'):
-			tk_addr = self.read_from_tmpfile('token_addr'+n).strip()
-			t = self.addrimport(
-				ext      = f'[{r}]{{}}.regtest.addrs',
-				expect   = expect,
-				add_args = ['--token-addr='+tk_addr]+extra_args )
-			t.p.wait()
-			ok_msg()
-		t.skip_ok = True
-		return t
+	def token_addrimport(self,addr_file,addr_range,expect,extra_args=[]):
+		token_addr = self.read_from_tmpfile(addr_file).strip()
+		return self.addrimport(
+			ext      = f'[{addr_range}]{{}}.regtest.addrs',
+			expect   = expect,
+			add_args = ['--token-addr='+token_addr]+extra_args )
+
+	def token_addrimport_addr1(self):
+		return self.token_addrimport('token_addr1','11-13',expect='3/3')
+
+	def token_addrimport_addr2(self):
+		return self.token_addrimport('token_addr2','21-23',expect='3/3')
 
 	def token_addrimport_batch(self):
-		return self.token_addrimport(extra_args=['--batch'],expect='OK: 3')
+		return self.token_addrimport('token_addr1','11-13',expect='OK: 3',extra_args=['--batch'])
+
+	def token_addrimport_sym(self):
+		return self.addrimport(
+			ext      = '[21-23]{}.regtest.addrs',
+			expect   = '3/3',
+			add_args = ['--token=MM2'] )
 
 	def bal7(self):       return self.bal5()
 	def token_bal1(self): return self.token_bal(n='1')