diff options
| author | Dana Powers <dana.powers@rd.io> | 2014-08-17 22:28:51 -0700 |
|---|---|---|
| committer | Dana Powers <dana.powers@rd.io> | 2014-08-19 21:27:43 -0700 |
| commit | 839b8831bff6c7a52112d0c413f8796617d7ee49 (patch) | |
| tree | 4ac0d285095ed4f69557f81d0a8928da262816db /test | |
| parent | 0a762a0bac61d83d686c062911a60cf11330d3a7 (diff) | |
| download | kafka-python-839b8831bff6c7a52112d0c413f8796617d7ee49.tar.gz | |
Cleanup some connection tests in test/test_conn.py
Diffstat (limited to 'test')
| -rw-r--r-- | test/test_conn.py | 25 |
1 files changed, 11 insertions, 14 deletions
diff --git a/test/test_conn.py b/test/test_conn.py index bcc6500..8534137 100644 --- a/test/test_conn.py +++ b/test/test_conn.py @@ -59,15 +59,11 @@ class ConnTest(unittest2.TestCase): 'payload': 'test data' } - def mock_reinit(obj): - obj._sock = mock.MagicMock() - obj._sock.sendall.return_value = None - obj._dirty = False - - with mock.patch.object(KafkaConnection, 'reinit', new=mock_reinit): - conn = KafkaConnection(fake_config['host'], fake_config['port']) - conn.send(fake_config['request_id'], fake_config['payload']) - conn._sock.sendall.assert_called_with(fake_config['payload']) + assert socket.create_connection is self.MockCreateConn + conn = KafkaConnection(fake_config['host'], fake_config['port']) + socket.create_connection.reset_mock() + conn.send(fake_config['request_id'], fake_config['payload']) + conn._sock.sendall.assert_called_with(fake_config['payload']) def test_init_creates_socket_connection(self): fake_config = { @@ -76,7 +72,8 @@ class ConnTest(unittest2.TestCase): } assert socket.create_connection is self.MockCreateConn - conn = KafkaConnection(fake_config['host'], fake_config['port']) + socket.create_connection.reset_mock() + KafkaConnection(fake_config['host'], fake_config['port']) socket.create_connection.assert_called_with((fake_config['host'], fake_config['port']), DEFAULT_SOCKET_TIMEOUT_SECONDS) def test_init_failure_raises_connection_error(self): @@ -88,9 +85,10 @@ class ConnTest(unittest2.TestCase): def raise_error(*args): raise socket.error - with mock.patch.object(socket, 'create_connection', new=raise_error): - with self.assertRaises(ConnectionError): - KafkaConnection(fake_config['host'], fake_config['port']) + assert socket.create_connection is self.MockCreateConn + socket.create_connection.side_effect=raise_error + with self.assertRaises(ConnectionError): + KafkaConnection(fake_config['host'], fake_config['port']) def test_send__reconnects_on_dirty_conn(self): fake_config = { @@ -105,7 +103,6 @@ class ConnTest(unittest2.TestCase): conn = KafkaConnection(fake_config['host'], fake_config['port']) # Dirty it - try: conn._raise_connection_error() except ConnectionError: |