# Path: xmanager/cloud/utils.py
# def get_master_address_port() -> Tuple[str, str]:
# def get_world_size_rank() -> Tuple[int, int]:
# def create_cluster_specs(workers: Sequence[str]) -> List[str]:
# def get_workerpool_address(workerpool: str) -> str:
# def map_workerpool_address_args(args: List[str]) -> List[str]:
# def print_workerpool_address_args(argv: List[str]) -> None:
# def create_workerpool_address_env_vars_script(path: str) -> None:
# def get_region() -> str:
# def get_closest_bucket(bucket_names: List[str]) -> str:


import os
import tempfile
import unittest
from xmanager.cloud import utils
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for xmanager.cloud.utils."""



_CLUSTER_SPEC = """{
  "cluster": {
    "workerpool0": ["cmle-training-workerpool0-ab-0:2222"],
    "workerpool1": ["cmle-training-workerpool1-ab-0:2222", "cmle-training-workerpool1-ab-1:2222"],
    "workerpool2": ["cmle-training-workerpool2-ab-0:2222", "cmle-training-workerpool2-ab-1:2222"]
  },
  "environment": "cloud",
  "task": {
    "type": "workerpool1",
    "index": 1,
    "trial": ""
  }
}""".replace('\n', ' ')


class UtilsTest(unittest.TestCase):

  def tearDown(self):
    super(UtilsTest, self).tearDown()
    os.environ['CLUSTER_SPEC'] = ''

  def test_get_master_address_port(self):
    os.environ['CLUSTER_SPEC'] = _CLUSTER_SPEC

# Path: xmanager/cloud/utils.py
# def get_master_address_port() -> Tuple[str, str]:
# def get_world_size_rank() -> Tuple[int, int]:
# def create_cluster_specs(workers: Sequence[str]) -> List[str]:
# def get_workerpool_address(workerpool: str) -> str:
# def map_workerpool_address_args(args: List[str]) -> List[str]:
# def print_workerpool_address_args(argv: List[str]) -> None:
# def create_workerpool_address_env_vars_script(path: str) -> None:
# def get_region() -> str:
# def get_closest_bucket(bucket_names: List[str]) -> str:


# Path: xmanager/cloud/utils_test.py
import os
import tempfile
import unittest
from xmanager.cloud import utils

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for xmanager.cloud.utils."""



_CLUSTER_SPEC = """{
  "cluster": {
    "workerpool0": ["cmle-training-workerpool0-ab-0:2222"],
    "workerpool1": ["cmle-training-workerpool1-ab-0:2222", "cmle-training-workerpool1-ab-1:2222"],
    "workerpool2": ["cmle-training-workerpool2-ab-0:2222", "cmle-training-workerpool2-ab-1:2222"]
  },
  "environment": "cloud",
  "task": {
    "type": "workerpool1",
    "index": 1,
    "trial": ""
  }
}""".replace('\n', ' ')


class UtilsTest(unittest.TestCase):

  def tearDown(self):
    super(UtilsTest, self).tearDown()
    os.environ['CLUSTER_SPEC'] = ''

  def test_get_master_address_port(self):
    os.environ['CLUSTER_SPEC'] = _CLUSTER_SPEC

    address, port = utils.get_master_address_port()