import os
import unittest

from core.websocket_client_v2 import QingjianAPIClient


class MultiWorkerRoutingTest(unittest.TestCase):
    def test_only_one_worker_owns_customer_when_no_explicit_shards(self):
        os.environ["AI_CS_WORKER_COUNT"] = "4"
        key = "shop_x:tb123456"
        owners = 0
        for wid in range(4):
            os.environ["AI_CS_WORKER_ID"] = str(wid)
            c = QingjianAPIClient(enable_agent=False)
            c.shard_keys = set()  # 模拟当前无分片表
            if c._is_owned_by_this_worker(key):
                owners += 1
        self.assertEqual(owners, 1)

    def tearDown(self):
        os.environ.pop("AI_CS_WORKER_COUNT", None)
        os.environ.pop("AI_CS_WORKER_ID", None)


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