From 8ad4b75806c0eb4f63cc1519993bcaba4fd18f8a Mon Sep 17 00:00:00 2001
From: huangyi <huang@crypto.com>
Date: Fri, 3 May 2024 15:36:48 +0800
Subject: [PATCH] fix test

---
 app/proposal.go                |  1 -
 integration_tests/cosmoscli.py |  4 +-
 integration_tests/test_e2ee.py | 89 ++++++++++++++++++++++++----------
 3 files changed, 65 insertions(+), 29 deletions(-)

diff --git a/app/proposal.go b/app/proposal.go
index 851ecb9a20..40646cd374 100644
--- a/app/proposal.go
+++ b/app/proposal.go
@@ -33,7 +33,6 @@ func NewProposalHandler(txDecoder sdk.TxDecoder, identity age.Identity) *Proposa
 }
 
 func (h *ProposalHandler) SetBlockList(blob []byte) error {
-	fmt.Println("SetBlockList")
 	if h.Identity == nil {
 		return nil
 	}
diff --git a/integration_tests/cosmoscli.py b/integration_tests/cosmoscli.py
index 11742dfdd3..af10bbab84 100644
--- a/integration_tests/cosmoscli.py
+++ b/integration_tests/cosmoscli.py
@@ -1857,7 +1857,7 @@ def query_e2ee_key(self, address):
                 home=self.data_dir,
                 output="json",
             )
-        )["key"]
+        ).get("key")
 
     def query_e2ee_keys(self, *addresses):
         return json.loads(
@@ -1869,7 +1869,7 @@ def query_e2ee_keys(self, *addresses):
                 home=self.data_dir,
                 output="json",
             )
-        )["keys"]
+        ).get("keys")
 
     def register_e2ee_key(self, key, **kwargs):
         kwargs.setdefault("gas_prices", DEFAULT_GAS_PRICE)
diff --git a/integration_tests/test_e2ee.py b/integration_tests/test_e2ee.py
index 9a8ae989f6..304ee3ee2f 100644
--- a/integration_tests/test_e2ee.py
+++ b/integration_tests/test_e2ee.py
@@ -1,40 +1,77 @@
+import json
+
+import pytest
+
+from .network import Cronos
+from .utils import wait_for_new_blocks
+
+
+def gen_validator_identity(cronos: Cronos):
+    for i in range(len(cronos.config["validators"])):
+        cli = cronos.cosmos_cli(i)
+        if cli.query_e2ee_key(cli.address("validator")):
+            return
+        pubkey = cli.keygen()
+        cli.register_e2ee_key(pubkey, _from="validator")
+        assert cli.query_e2ee_key(cli.address("validator")) == pubkey
+
+        cronos.supervisorctl("restart", f"cronos_777-1-node{i}")
+
+    wait_for_new_blocks(cronos.cosmos_cli(), 1)
+
+
 def test_encrypt_decrypt(cronos):
+    gen_validator_identity(cronos)
+
     cli0 = cronos.cosmos_cli()
     cli1 = cronos.cosmos_cli(1)
 
-    # gen two keys for two accounts
-    name0 = "key0"
-    name1 = "key1"
-    pubkey0 = cli0.keygen(keyring_name=name0)
-    pubkey1 = cli1.keygen(keyring_name=name1)
-    sender = "validator"
-    cli0.register_e2ee_key(pubkey0, _from=sender)
-    cli1.register_e2ee_key(pubkey1, _from=sender)
     # query in batch
-    assert cli0.query_e2ee_keys(cli0.address(sender), cli1.address(sender)) == [
-        pubkey0,
-        pubkey1,
-    ]
+    assert (
+        len(
+            cli0.query_e2ee_keys(
+                cli0.address("validator"),
+                cli1.address("validator"),
+            )
+        )
+        == 2
+    )
+
     # prepare data file to encrypt
     content = "Hello World!"
-def test_block_list(cronos):
-    cli0 = cronos.cosmos_cli()
-    cli1 = cronos.cosmos_cli(1)
+    plainfile = cli0.data_dir / "plaintext"
+    plainfile.write_text(content)
+    cipherfile = cli0.data_dir / "ciphertext"
+    cli0.encrypt(
+        plainfile,
+        cli0.address("validator"),
+        cli1.address("validator"),
+        output=cipherfile,
+    )
+
+    assert cli0.decrypt(cipherfile) == content
+    assert cli1.decrypt(cipherfile) == content
 
-    # prepare encryption keys for validators
-    cli0.register_e2ee_key(cli0.keygen(), _from="validator")
-    cli1.register_e2ee_key(cli1.keygen(), _from="validator")
 
-    user = cli0.address("user")
+def test_block_list(cronos):
+    gen_validator_identity(cronos)
+    cli = cronos.cosmos_cli()
+
+    user = cli.address("user")
 
     blocklist = json.dumps({"addresses": [user]})
-    plainfile = cli0.data_dir / "plaintext"
+    plainfile = cli.data_dir / "plaintext"
     plainfile.write_text(blocklist)
-    cipherfile = cli0.data_dir / "ciphertext"
-    cli0.encrypt_to_validators(plainfile, output=cipherfile)
-    rsp = cli0.store_blocklist(cipherfile, _from="validator")
+    cipherfile = cli.data_dir / "ciphertext"
+    cli.encrypt_to_validators(plainfile, output=cipherfile)
+    rsp = cli.store_blocklist(cipherfile, _from="validator")
     assert rsp["code"] == 0, rsp["raw_log"]
 
-    wait_for_new_blocks(cli0, 2)
-    rsp = cli0.transfer(user, cli0.address("validator"), "1basetcro")
-    assert rsp["code"] != 0
+    # normal tx works
+    cli.transfer(cli.address("validator"), user, "1basetcro")
+
+    # blocked tx don't work
+    with pytest.raises(AssertionError) as exc:
+        cli.transfer(user, cli.address("validator"), "1basetcro")
+
+    assert "timed out waiting for event" in str(exc.value)