def test_generate_with_dhe(self):
state = ConnectionState()
state.key_exchange = mock.MagicMock()
cke = ClientKeyExchangeGenerator(
cipher=constants.CipherSuite.TLS_DHE_RSA_WITH_AES_128_CBC_SHA)
ret = cke.generate(state)
self.assertIs(ret, state.key_exchange.makeClientKeyExchange())
