Mocking socket.recv() - python

I want to mock the recv function of a socket and having trouble. Getting errors that I need to actually connect my socket even though I should not have to.
Here is all relevant code:
Class A:
def __init__.py(self):
self.s = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
def connect(self):
self.s.connect((host,port))
def data_collection(self, d):
self.send_over_socket(self.cg.get_privilege_command())
self.send_over_socket(self.status_string)
buf = ""
while 1:
buf = buf + self.s.recv(self.BUFFER_SIZE).decode()
#Removed string manipulation below
self.publish(data from string manipulation)
def send_over_socket(self,data):
return (self.s.send(bytes(data.encode())))
Pytest code;
def test_data_collection():
with mock.patch('moduleA.classA.send_over_socket') as mock_send_over_socket:
A = classA()
A.send_over_socket.return_value = 1
with mock.patch('socket.socket') as mock_socket:
mock_socket.return_value.recv.decode.return_value = "packet_string"
with mock.patch('moduleA.ClassA.publish') as mock_publish:
d = {}
A.data_collection(d)
A.publish.assert_called_with("0","500","2","1")

You're missing a return_value for your mock_socket; you need to work off of what socket.socket is going to return. Your test should look like this:
def test_data_collection():
with mock.patch('socket.socket') as mock_socket:
A = classA()
mock_socket.return_value.recv.decode.return_value = "ABC123"
buf = A.data_collection()
assert "ABC123" == buf
Note that I had to fix some syntax errors in your class code, and I made data_collection return the buf variable so that it can be tested.
EDIT: File with class:
"""receiver.py"""
import socket
IP = 'localhost'
PORT = 80
class classA:
def __init__(self):
self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
def connect(self):
self.s.connect((IP, PORT))
def data_collection(self):
buf = self.s.recv.decode()
return buf

Related

How I can check connection correctly

I have a boolean variable:
is_conected = False
And I have this part of code:
def start(self, host='localhost', port=32767):
self.__connection.bind((host, port))
self.__connection.listen(5)
def send_frame(self, frame: str):
if not self.__connection.send(bytes(frame, encoding='utf-8')):
raise RuntimeError("Connection broken.")
def recieve_frame(self) -> str:
socket, adress = self.__connection.accept()
str(__data) = socket.recv(1024)
__frame_buffer.append(__data)
Where I need to put If... Else, for checking, if someone connected or not? And what kind of condition I need to use?
The whole code:
import socket
class TCPServer():
__frame_buffer = list()
__connection = None
is_conected = True
def __new__(cls):
if not hasattr(cls, 'instance'):
cls.instance = super(TCPServer, cls).__new__(cls)
return cls.instance
def __init__(self):
self.__connection = socket.socket(
socket.AF_INET, socket.SOCK_STREAM)
def start(self, host='localhost', port=32767):
# """ Start server """
self.__connection.bind((host, port))
self.__connection.listen(5)
def send_frame(self, frame: str):
# """ Send frame """
if not self.__connection.send(bytes(frame, encoding='utf-8')):
raise RuntimeError("Connection broken.")
def recieve_frame(self) -> str:
socket, adress = self.__connection.accept()
str(__data) = socket.recv(1024)
__frame_buffer.append(__data)
def get_frame(self) -> str:
if self.__frame_buffer:
self.__frame_buffer.pop(0)
server = TCPServer()
while True:
if server.is_conected:
server.send_frame(str(input("String to send")))
print(server.get_frame())
That's the whole code. I don't write a more part of it, only a send and recieve frame function. So, that's kinda a school project for me and I'm not too good in working with any kind of connection or protocols
This is how the server generally works
with socket(AF_INET, SOCK_STREAM) as s:
s.bind((HOST, PORT))
s.listen() # Start listening
while True:
conn, addr = s.accept() # Accept connection request
Connected = True # You are connected
try:
#Work with connection
except:
Connected = False
Figure out how you can do this with your code

Mocked method socket.accept() is not called

I am trying to mock server side of socket. This example is not working and I don't know why (it still uses original socket.socket.accept() method).
Any ideas how to mock listening socket correctly?
class MySocket:
def __init__:
self.socket = socket.socket()
def listen(self):
self.socket.accept() # original socket method
class MockConn:
def recv(self, int):
return b''
def close(self):
pass
def sendall(self, str):
pass
with mock.patch('socket.socket') as mock_socket:
conn = MockConn()
mock_socket.return_value.accept.return_value = (conn, 'localhost')
s = MySocket(None)
s.listen() # this should call accept from the mock
Not sure what you want, but this works:
import socket
from unittest import mock
class MySocket:
def listen(self):
print('Accept result:', socket.socket().accept())
# Accept result: (<__main__.MockConn object at 0x000001E592B36B38>, 'localhost')
class MockConn:
def recv(self, int):
return b''
def close(self):
pass
def sendall(self, str):
pass
with mock.patch('socket.socket') as mock_socket:
conn = MockConn()
mock_socket.return_value.accept.return_value = (conn, 'localhost')
s = MySocket()
s.listen()
With some very minor changes to your current code this works:
from unittest import mock
import socket
class MySocket:
def __init__(self): # init takes self
self.socket = socket.socket()
def listen(self):
return self.socket.accept() # original socket method
class MockConn:
def recv(self, int):
return b''
def close(self):
pass
def sendall(self, str):
pass
with mock.patch('socket.socket') as mock_socket:
conn = MockConn()
mock_socket.return_value.accept.return_value = (conn, 'localhost')
s = MySocket()
print(s.listen()) # this should call accept from the mock
Output:
(<__main__.MockConn object at 0x00000255C5A53B00>, 'localhost')

Python Socketserver client timeout

I have a socketserver in Python which has to handle multiple clients using the select.select method, as seen in the code below:
import socket
import select
class Server:
def __init__(self):
self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.server.bind(('localhost', 2000))
self.socket_list = [self.server]
self.addresses = ['localhost']
self.commands = [""]
self.running = True
self.server.listen(10)
def listen(self):
while self.running:
read, write, error = select.select(self.socket_list, [], self.socket_list, 0)
for sock in read:
if sock == self.server and self.running:
try:
conn, address = self.server.accept()
conn.settimeout(30)
self.socket_list.append(conn)
self.addresses.append(address[0])
self.commands.append("")
except:
self.shutdown()
break
elif self.running:
try:
packet = sock.recv(60)
if not packet:
self.close_conn(sock)
index = self.socket_list.index(sock)
self.commands[index] += packet
if '\n' in self.commands[index]:
#handle command
except:
self.close_conn(sock)
def close_conn(self, conn):
#close client conn
def shutdown(self):
#shutdown server
if __name__ == "__main__":
Server().listen()
The problem i currently have is that the client's connection should be closed after 30 seconds, but that doesn't happen, even though i declared that by using conn.settimeout(30). I haven't found any explanation yet as to why this happens.
Note: comments were used to replace parts of the code that didn't mater to the problem.

How to make a dynamic port forwarding on python?

I modified the code (found here) a bit (added class AsyncEventLoop)
import socket,asyncore
import threading
class forwarder(asyncore.dispatcher):
def __init__(self, ip, port, remoteip,remoteport,backlog=5):
asyncore.dispatcher.__init__(self)
self.remoteip=remoteip
self.remoteport=remoteport
self.create_socket(socket.AF_INET,socket.SOCK_STREAM)
self.set_reuse_addr()
self.bind((ip,port))
self.listen(backlog)
def handle_accept(self):
conn, addr = self.accept()
# print '--- Connect --- '
sender(receiver(conn),self.remoteip,self.remoteport)
def closef(self):
self.close()
class receiver(asyncore.dispatcher):
def __init__(self,conn):
asyncore.dispatcher.__init__(self,conn)
self.from_remote_buffer=''
self.to_remote_buffer=''
self.sender=None
def handle_connect(self):
pass
def handle_read(self):
read = self.recv(4096)
# print '%04i -->'%len(read)
self.from_remote_buffer += read
def writable(self):
return (len(self.to_remote_buffer) > 0)
def handle_write(self):
sent = self.send(self.to_remote_buffer)
# print '%04i <--'%sent
self.to_remote_buffer = self.to_remote_buffer[sent:]
def handle_close(self):
self.close()
if self.sender:
self.sender.close()
class sender(asyncore.dispatcher):
def __init__(self, receiver, remoteaddr,remoteport):
asyncore.dispatcher.__init__(self)
self.receiver=receiver
receiver.sender=self
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.connect((remoteaddr, remoteport))
def handle_connect(self):
pass
def handle_read(self):
read = self.recv(4096)
# print '<-- %04i'%len(read)
self.receiver.to_remote_buffer += read
def writable(self):
return (len(self.receiver.from_remote_buffer) > 0)
def handle_write(self):
sent = self.send(self.receiver.from_remote_buffer)
# print '--> %04i'%sent
self.receiver.from_remote_buffer = self.receiver.from_remote_buffer[sent:]
def handle_close(self):
self.close()
self.receiver.close()
class AsyncEventLoop(threading.Thread):
def run(self):
asyncore.loop()
starting in the code
ser = forwarder('127.0.0.1', 7774, 'google.com.ua', 80)
evLoop = AsyncEventLoop()
Can I change remoteip, remoteport on the fly? If the fly is not possible to change, as well stop this thread to start with the new settings (remoteip, remoteport)?
yes, you can
ser.remoteip, ser.remoteport = IP, port

How do i control what socket thread I want to connect to in a asychronous reverse server python?

Good evening, This is my 1st time on this site, I have been programming a python based user monitoring system for my work for the past 3 months and I am almost done with my 1st release. However I have run into a problem controlling what computer I want to connect to.
If i run the two sample code I put in this post I can receive the client and send commands to client with the server, but only one client at a time, and the server is dictating which client I can send to and which one is next. I am certain the problem is "server side but I am not sure how to fix the problem and a Google search does not turn up anyone having tried this.
I have attached both client and server base networking code in this post.
client:
import asyncore
import socket
import sys
do_restart = False
class client(asyncore.dispatcher):
def __init__(self, host, port=8000):
serv = open("srv.conf","r")
host = serv.read()
serv.close()
asyncore.dispatcher.__init__(self)
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.connect((host, port))
def writable(self):
return 0
def handle_connect(self):
pass
def handle_read(self):
data = self.recv(4096)
#Rest of code goes here
serv = open("srv.conf","r")
host = serv.read()
serv.close()
request = client(host)
asyncore.loop()
server:
import asyncore
import socket
import sys
class soc(asyncore.dispatcher):
def __init__(self, port=8000):
asyncore.dispatcher.__init__(self)
self.port = port
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.bind(('', port))
self.listen(5)
def handle_accept(self):
channel, addr = self.accept()
while 1:
j = raw_input(addr)
#Rest of my code is here
server = soc(8000)
asyncore.loop()
Here is a fast and dirty idea that I threw together.
The use of raw_input has been replaced with another dispatcher that is asyncore compatable, referencing this other question here
And I am expanding on the answer given by #user1320237 to defer each new connection to a new dispatcher.
You wanted to have a single command line interface that can send control commands to any of the connected clients. That means you need a way to switch between them. What I have done is created a dict to keep track of the connected clients. Then we also create a set of available commands that map to callbacks for your command line.
This example has the following:
list: list current clients
set <client>: set current client
send <msg>: send a msg to the current client
server.py
import asyncore
import socket
import sys
from weakref import WeakValueDictionary
class Soc(asyncore.dispatcher):
CMDS = {
'list': 'cmd_list',
'set': 'cmd_set_addr',
'send': 'cmd_send',
}
def __init__(self, port=8000):
asyncore.dispatcher.__init__(self)
self._conns = WeakValueDictionary()
self._current = tuple()
self.port = port
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.set_reuse_addr()
self.bind(('', port))
self.listen(5)
self.cmdline = Cmdline(self.handle_input, sys.stdin)
self.cmdline.prompt()
def writable(self):
return False
def handle_input(self, i):
tokens = i.strip().split(None, 1)
cmd = tokens[0]
arg = ""
if len(tokens) > 1:
arg = tokens[1]
cbk = self.CMDS.get(cmd)
if cbk:
getattr(self, cbk)(arg)
self.cmdline.prompt(self._addr_to_key(self._current))
def handle_accept(self):
channel, addr = self.accept()
c = Conn(channel)
self._conns[self._addr_to_key(addr)] = c
def _addr_to_key(self, addr):
return ':'.join(str(i) for i in addr)
def cmd_list(self, *args):
avail = '\n'.join(self._conns.iterkeys())
print "\n%s\n" % avail
def cmd_set_addr(self, addr_str):
conn = self._conns.get(addr_str)
if conn:
self._current = conn.addr
def cmd_send(self, msg):
if self._current:
addr_str = self._addr_to_key(self._current)
conn = self._conns.get(addr_str)
if conn:
conn.buffer += msg
class Cmdline(asyncore.file_dispatcher):
def __init__(self, cbk, f):
asyncore.file_dispatcher.__init__(self, f)
self.cbk = cbk
def prompt(self, msg=''):
sys.stdout.write('%s > ' % msg)
sys.stdout.flush()
def handle_read(self):
self.cbk(self.recv(1024))
class Conn(asyncore.dispatcher):
def __init__(self, *args, **kwargs):
asyncore.dispatcher.__init__(self, *args, **kwargs)
self.buffer = ""
def writable(self):
return len(self.buffer) > 0
def handle_write(self):
self.send(self.buffer)
self.buffer = ''
def handle_read(self):
data = self.recv(4096)
print self.addr, '-', data
server = Soc(8000)
asyncore.loop()
Your main server is now never blocking on stdin, and always accepting new connections. The only work it does is the command handling which should either be a fast operation, or signals the connection objects to handle the message.
Usage:
# start the server
# start 2 clients
>
> list
127.0.0.1:51738
127.0.0.1:51736
> set 127.0.0.1:51736
127.0.0.1:51736 >
127.0.0.1:51736 > send foo
# client 127.0.0.1:51736 receives "foo"
To me
while 1:
j = raw_input(addr)
seems to be the problem:
you only accept a socket an then do something with it until end.
You should create e new dispatcher for every client connecting
class conn(asyncore.dispatcher):
...
def handle_read(self):
...
class soc(asyncore.dispatcher):
def handle_accept(self):
...
c = conn()
c.set_socket(channel)
Asyncore will call you back for every read operation possible.
Asyncore uses only one thread. This is its strength. every dispatcher that has a socket is called one after an other with those handle_* functions.

Categories

Resources