def test_no_factory_method_found(self):
factory = RendezvousHandlerFactory()
rdzv_params = RendezvousParameters(
backend="mock", endpoint="", run_id="foobar", min_nodes=1, max_nodes=2
)
with self.assertRaises(ValueError):
factory.create_rdzv_handler(rdzv_params)
def test_create_rdzv_handler(self):
rdzv_params = RendezvousParameters(
backend="mock", endpoint="", run_id="foobar", min_nodes=1, max_nodes=2
)
factory = RendezvousHandlerFactory()
factory.register("mock", create_mock_rdzv_handler)
mock_rdzv_handler = factory.create_rdzv_handler(rdzv_params)
self.assertTrue(isinstance(mock_rdzv_handler, MockRendezvousHandler))
def test_double_registration(self):
factory = RendezvousHandlerFactory()
factory.register("mock", create_mock_rdzv_handler)
with self.assertRaises(ValueError):
factory.register("mock", create_mock_rdzv_handler)
def get_rendezvous_handler(rdzv_params: RendezvousParameters) -> RendezvousHandler:
factory = RendezvousHandlerFactory()
factory.register("etcd", etcd_rdzv.create_rdzv_handler)
return factory.create_rdzv_handler(rdzv_params)