import unittest
from bagua.torch_api.contrib.utils.redis_store import (
    RedisStore,
    start_redis_server_cli,
)
from bagua.torch_api.env import find_free_network_port
import redis
import multiprocessing as mp
import logging
import numpy as np
import pickle
from tests import skip_if_cuda_available


logging.basicConfig(level=logging.DEBUG)


class TestRedisStore(unittest.TestCase):
    def check(self, store):
        store.clear()
        self.assertEqual(store.num_keys(), 0)

        self.generated_data = [np.random.rand(10) for _ in range(5)]
        store.set("1", pickle.dumps(self.generated_data[1]))

        store.mset(
            {
                "2": pickle.dumps(self.generated_data[2]),
                "3": pickle.dumps(self.generated_data[3]),
                "4": pickle.dumps(self.generated_data[4]),
            }
        )
        ret = store.mget(["1", "2", "5"])
        self.assertTrue((pickle.loads(ret[0]) == self.generated_data[1]).all())
        self.assertTrue((pickle.loads(ret[1]) == self.generated_data[2]).all())
        self.assertEqual(ret[2], None)

        r1 = store.get("4")
        r2 = store.get("6")
        self.assertTrue((pickle.loads(r1) == self.generated_data[4]).all())
        self.assertEqual(r2, None)

        cnt = store.num_keys()
        self.assertEqual(cnt, 4)

        self.assertTrue(store.status())

    @skip_if_cuda_available()
    def test_redis_store(self):
        store = RedisStore(hosts=None, cluster_mode=False, capacity_per_node=10000000)
        self.check(store)

    @skip_if_cuda_available()
    def test_redis_cluster_store(self):
        n = 3
        hosts = []
        ports = []
        processes = []
        for i in range(n):
            port = find_free_network_port()
            p = mp.Process(
                target=start_redis_server_cli,
                args=(port, 10000000, f"--cluster-config-file nodes{port}.conf"),
            )
            p.start()

            ports.append(port)
            processes.append(p)
            hosts.append({"host": "127.0.0.1", "port": port})

        for p in processes:
            p.join()
            self.assertTrue(p.exitcode == 0)

        store = RedisStore(hosts=hosts, cluster_mode=True)
        self.check(store)

        # Shut down servers manually
        for port in ports:
            client = redis.Redis(port=port)
            client.shutdown(nosave=True)
            client.close()


if __name__ == "__main__":
    unittest.main()
