Mercurial > public > mercurial-scm > hg
comparison mercurial/util.py @ 37010:8453699a1f21
util: observable proxy objects for sockets
We previously introduced proxy objects and observers for file objects
to help implement low-level tests for the SSH wire protocol.
In this commit, we do the same for sockets in order to help test the
HTTP server.
We only proxy/observe some socket methods. I didn't feel like
implementing all the methods because there are so many of them and
implementing them will provide no short term value. We can always
implement them later.
# no-check-commit because we implement foo_bar methods on stdlib types
Differential Revision: https://phab.mercurial-scm.org/D2721
author | Gregory Szorc <gregory.szorc@gmail.com> |
---|---|
date | Tue, 13 Mar 2018 17:42:00 -0700 |
parents | d683c7367989 |
children | d3a9036d9ae9 |
comparison
equal
deleted
inserted
replaced
37009:5890e5872f36 | 37010:8453699a1f21 |
---|---|
687 if fn: | 687 if fn: |
688 fn(res) | 688 fn(res) |
689 | 689 |
690 return res | 690 return res |
691 | 691 |
692 PROXIED_SOCKET_METHODS = { | |
693 r'makefile', | |
694 r'recv', | |
695 r'recvfrom', | |
696 r'recvfrom_into', | |
697 r'recv_into', | |
698 r'send', | |
699 r'sendall', | |
700 r'sendto', | |
701 r'setblocking', | |
702 r'settimeout', | |
703 r'gettimeout', | |
704 r'setsockopt', | |
705 } | |
706 | |
707 class socketproxy(object): | |
708 """A proxy around a socket that tells a watcher when events occur. | |
709 | |
710 This is like ``fileobjectproxy`` except for sockets. | |
711 | |
712 This type is intended to only be used for testing purposes. Think hard | |
713 before using it in important code. | |
714 """ | |
715 __slots__ = ( | |
716 r'_orig', | |
717 r'_observer', | |
718 ) | |
719 | |
720 def __init__(self, sock, observer): | |
721 object.__setattr__(self, r'_orig', sock) | |
722 object.__setattr__(self, r'_observer', observer) | |
723 | |
724 def __getattribute__(self, name): | |
725 if name in PROXIED_SOCKET_METHODS: | |
726 return object.__getattribute__(self, name) | |
727 | |
728 return getattr(object.__getattribute__(self, r'_orig'), name) | |
729 | |
730 def __delattr__(self, name): | |
731 return delattr(object.__getattribute__(self, r'_orig'), name) | |
732 | |
733 def __setattr__(self, name, value): | |
734 return setattr(object.__getattribute__(self, r'_orig'), name, value) | |
735 | |
736 def __nonzero__(self): | |
737 return bool(object.__getattribute__(self, r'_orig')) | |
738 | |
739 __bool__ = __nonzero__ | |
740 | |
741 def _observedcall(self, name, *args, **kwargs): | |
742 # Call the original object. | |
743 orig = object.__getattribute__(self, r'_orig') | |
744 res = getattr(orig, name)(*args, **kwargs) | |
745 | |
746 # Call a method on the observer of the same name with arguments | |
747 # so it can react, log, etc. | |
748 observer = object.__getattribute__(self, r'_observer') | |
749 fn = getattr(observer, name, None) | |
750 if fn: | |
751 fn(res, *args, **kwargs) | |
752 | |
753 return res | |
754 | |
755 def makefile(self, *args, **kwargs): | |
756 res = object.__getattribute__(self, r'_observedcall')( | |
757 r'makefile', *args, **kwargs) | |
758 | |
759 # The file object may be used for I/O. So we turn it into a | |
760 # proxy using our observer. | |
761 observer = object.__getattribute__(self, r'_observer') | |
762 return makeloggingfileobject(observer.fh, res, observer.name, | |
763 reads=observer.reads, | |
764 writes=observer.writes, | |
765 logdata=observer.logdata) | |
766 | |
767 def recv(self, *args, **kwargs): | |
768 return object.__getattribute__(self, r'_observedcall')( | |
769 r'recv', *args, **kwargs) | |
770 | |
771 def recvfrom(self, *args, **kwargs): | |
772 return object.__getattribute__(self, r'_observedcall')( | |
773 r'recvfrom', *args, **kwargs) | |
774 | |
775 def recvfrom_into(self, *args, **kwargs): | |
776 return object.__getattribute__(self, r'_observedcall')( | |
777 r'recvfrom_into', *args, **kwargs) | |
778 | |
779 def recv_into(self, *args, **kwargs): | |
780 return object.__getattribute__(self, r'_observedcall')( | |
781 r'recv_info', *args, **kwargs) | |
782 | |
783 def send(self, *args, **kwargs): | |
784 return object.__getattribute__(self, r'_observedcall')( | |
785 r'send', *args, **kwargs) | |
786 | |
787 def sendall(self, *args, **kwargs): | |
788 return object.__getattribute__(self, r'_observedcall')( | |
789 r'sendall', *args, **kwargs) | |
790 | |
791 def sendto(self, *args, **kwargs): | |
792 return object.__getattribute__(self, r'_observedcall')( | |
793 r'sendto', *args, **kwargs) | |
794 | |
795 def setblocking(self, *args, **kwargs): | |
796 return object.__getattribute__(self, r'_observedcall')( | |
797 r'setblocking', *args, **kwargs) | |
798 | |
799 def settimeout(self, *args, **kwargs): | |
800 return object.__getattribute__(self, r'_observedcall')( | |
801 r'settimeout', *args, **kwargs) | |
802 | |
803 def gettimeout(self, *args, **kwargs): | |
804 return object.__getattribute__(self, r'_observedcall')( | |
805 r'gettimeout', *args, **kwargs) | |
806 | |
807 def setsockopt(self, *args, **kwargs): | |
808 return object.__getattribute__(self, r'_observedcall')( | |
809 r'setsockopt', *args, **kwargs) | |
810 | |
692 DATA_ESCAPE_MAP = {pycompat.bytechr(i): br'\x%02x' % i for i in range(256)} | 811 DATA_ESCAPE_MAP = {pycompat.bytechr(i): br'\x%02x' % i for i in range(256)} |
693 DATA_ESCAPE_MAP.update({ | 812 DATA_ESCAPE_MAP.update({ |
694 b'\\': b'\\\\', | 813 b'\\': b'\\\\', |
695 b'\r': br'\r', | 814 b'\r': br'\r', |
696 b'\n': br'\n', | 815 b'\n': br'\n', |
701 if isinstance(s, bytearray): | 820 if isinstance(s, bytearray): |
702 s = bytes(s) | 821 s = bytes(s) |
703 | 822 |
704 return DATA_ESCAPE_RE.sub(lambda m: DATA_ESCAPE_MAP[m.group(0)], s) | 823 return DATA_ESCAPE_RE.sub(lambda m: DATA_ESCAPE_MAP[m.group(0)], s) |
705 | 824 |
706 class fileobjectobserver(object): | 825 class baseproxyobserver(object): |
826 def _writedata(self, data): | |
827 if not self.logdata: | |
828 self.fh.write('\n') | |
829 self.fh.flush() | |
830 return | |
831 | |
832 # Simple case writes all data on a single line. | |
833 if b'\n' not in data: | |
834 self.fh.write(': %s\n' % escapedata(data)) | |
835 self.fh.flush() | |
836 return | |
837 | |
838 # Data with newlines is written to multiple lines. | |
839 self.fh.write(':\n') | |
840 lines = data.splitlines(True) | |
841 for line in lines: | |
842 self.fh.write('%s> %s\n' % (self.name, escapedata(line))) | |
843 self.fh.flush() | |
844 | |
845 class fileobjectobserver(baseproxyobserver): | |
707 """Logs file object activity.""" | 846 """Logs file object activity.""" |
708 def __init__(self, fh, name, reads=True, writes=True, logdata=False): | 847 def __init__(self, fh, name, reads=True, writes=True, logdata=False): |
709 self.fh = fh | 848 self.fh = fh |
710 self.name = name | 849 self.name = name |
711 self.logdata = logdata | 850 self.logdata = logdata |
712 self.reads = reads | 851 self.reads = reads |
713 self.writes = writes | 852 self.writes = writes |
714 | 853 |
715 def _writedata(self, data): | |
716 if not self.logdata: | |
717 self.fh.write('\n') | |
718 self.fh.flush() | |
719 return | |
720 | |
721 # Simple case writes all data on a single line. | |
722 if b'\n' not in data: | |
723 self.fh.write(': %s\n' % escapedata(data)) | |
724 self.fh.flush() | |
725 return | |
726 | |
727 # Data with newlines is written to multiple lines. | |
728 self.fh.write(':\n') | |
729 lines = data.splitlines(True) | |
730 for line in lines: | |
731 self.fh.write('%s> %s\n' % (self.name, escapedata(line))) | |
732 self.fh.flush() | |
733 | |
734 def read(self, res, size=-1): | 854 def read(self, res, size=-1): |
735 if not self.reads: | 855 if not self.reads: |
736 return | 856 return |
737 # Python 3 can return None from reads at EOF instead of empty strings. | 857 # Python 3 can return None from reads at EOF instead of empty strings. |
738 if res is None: | 858 if res is None: |
790 """Turn a file object into a logging file object.""" | 910 """Turn a file object into a logging file object.""" |
791 | 911 |
792 observer = fileobjectobserver(logh, name, reads=reads, writes=writes, | 912 observer = fileobjectobserver(logh, name, reads=reads, writes=writes, |
793 logdata=logdata) | 913 logdata=logdata) |
794 return fileobjectproxy(fh, observer) | 914 return fileobjectproxy(fh, observer) |
915 | |
916 class socketobserver(baseproxyobserver): | |
917 """Logs socket activity.""" | |
918 def __init__(self, fh, name, reads=True, writes=True, states=True, | |
919 logdata=False): | |
920 self.fh = fh | |
921 self.name = name | |
922 self.reads = reads | |
923 self.writes = writes | |
924 self.states = states | |
925 self.logdata = logdata | |
926 | |
927 def makefile(self, res, mode=None, bufsize=None): | |
928 if not self.states: | |
929 return | |
930 | |
931 self.fh.write('%s> makefile(%r, %r)\n' % ( | |
932 self.name, mode, bufsize)) | |
933 | |
934 def recv(self, res, size, flags=0): | |
935 if not self.reads: | |
936 return | |
937 | |
938 self.fh.write('%s> recv(%d, %d) -> %d' % ( | |
939 self.name, size, flags, len(res))) | |
940 self._writedata(res) | |
941 | |
942 def recvfrom(self, res, size, flags=0): | |
943 if not self.reads: | |
944 return | |
945 | |
946 self.fh.write('%s> recvfrom(%d, %d) -> %d' % ( | |
947 self.name, size, flags, len(res[0]))) | |
948 self._writedata(res[0]) | |
949 | |
950 def recvfrom_into(self, res, buf, size, flags=0): | |
951 if not self.reads: | |
952 return | |
953 | |
954 self.fh.write('%s> recvfrom_into(%d, %d) -> %d' % ( | |
955 self.name, size, flags, res[0])) | |
956 self._writedata(buf[0:res[0]]) | |
957 | |
958 def recv_into(self, res, buf, size=0, flags=0): | |
959 if not self.reads: | |
960 return | |
961 | |
962 self.fh.write('%s> recv_into(%d, %d) -> %d' % ( | |
963 self.name, size, flags, res)) | |
964 self._writedata(buf[0:res]) | |
965 | |
966 def send(self, res, data, flags=0): | |
967 if not self.writes: | |
968 return | |
969 | |
970 self.fh.write('%s> send(%d, %d) -> %d' % ( | |
971 self.name, len(data), flags, len(res))) | |
972 self._writedata(data) | |
973 | |
974 def sendall(self, res, data, flags=0): | |
975 if not self.writes: | |
976 return | |
977 | |
978 # Returns None on success. So don't bother reporting return value. | |
979 self.fh.write('%s> sendall(%d, %d)' % ( | |
980 self.name, len(data), flags)) | |
981 self._writedata(data) | |
982 | |
983 def sendto(self, res, data, flagsoraddress, address=None): | |
984 if not self.writes: | |
985 return | |
986 | |
987 if address: | |
988 flags = flagsoraddress | |
989 else: | |
990 flags = 0 | |
991 | |
992 self.fh.write('%s> sendto(%d, %d, %r) -> %d' % ( | |
993 self.name, len(data), flags, address, res)) | |
994 self._writedata(data) | |
995 | |
996 def setblocking(self, res, flag): | |
997 if not self.states: | |
998 return | |
999 | |
1000 self.fh.write('%s> setblocking(%r)\n' % (self.name, flag)) | |
1001 | |
1002 def settimeout(self, res, value): | |
1003 if not self.states: | |
1004 return | |
1005 | |
1006 self.fh.write('%s> settimeout(%r)\n' % (self.name, value)) | |
1007 | |
1008 def gettimeout(self, res): | |
1009 if not self.states: | |
1010 return | |
1011 | |
1012 self.fh.write('%s> gettimeout() -> %f\n' % (self.name, res)) | |
1013 | |
1014 def setsockopt(self, level, optname, value): | |
1015 if not self.states: | |
1016 return | |
1017 | |
1018 self.fh.write('%s> setsockopt(%r, %r, %r) -> %r\n' % ( | |
1019 self.name, level, optname, value)) | |
1020 | |
1021 def makeloggingsocket(logh, fh, name, reads=True, writes=True, states=True, | |
1022 logdata=False): | |
1023 """Turn a socket into a logging socket.""" | |
1024 | |
1025 observer = socketobserver(logh, name, reads=reads, writes=writes, | |
1026 states=states, logdata=logdata) | |
1027 return socketproxy(fh, observer) | |
795 | 1028 |
796 def version(): | 1029 def version(): |
797 """Return version information if available.""" | 1030 """Return version information if available.""" |
798 try: | 1031 try: |
799 from . import __version__ | 1032 from . import __version__ |