Skip to content

Commit

Permalink
Fix mock SonicV2Connector in python3: use decode_responses mode so ca…
Browse files Browse the repository at this point in the history
…ller code will be the same as python2 (sonic-net#1238)

The python2 behavior is not changed.
  • Loading branch information
qiluo-msft authored Nov 14, 2020
1 parent 71aa221 commit 8079558
Showing 1 changed file with 28 additions and 9 deletions.
37 changes: 28 additions & 9 deletions tests/mock_tables/dbconnector.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# MONKEY PATCH!!!
import json
import os
import sys

import mock
import mockredis
Expand All @@ -10,6 +11,13 @@
from swsssdk import SonicDBConfig, SonicV2Connector
from swsscommon import swsscommon


if sys.version_info >= (3, 0):
long = int
xrange = range
basestring = str
from functools import reduce

topo = None

def clean_up_config():
Expand Down Expand Up @@ -46,6 +54,7 @@ def connect_SonicV2Connector(self, db_name, retry_on=True):
self.dbintf.redis_kwargs['namespace'] = self.namespace
# Mock DB filename for unit-test
self.dbintf.redis_kwargs['db_name'] = db_name
self.dbintf.redis_kwargs['decode_responses'] = True
_old_connect_SonicV2Connector(self, db_name, retry_on)

def _subscribe_keyspace_notification(self, db_name, client):
Expand Down Expand Up @@ -86,6 +95,7 @@ def __init__(self, *args, **kwargs):
topo = kwargs.pop('topo')
namespace = kwargs.pop('namespace')
db_name = kwargs.pop('db_name')
self.decode_responses = kwargs.pop('decode_responses', False) == True
fname = db_name.lower() + ".json"
self.pubsub = MockPubSub()

Expand All @@ -110,6 +120,23 @@ def __init__(self, *args, **kwargs):
for attr, value in v.items():
self.hset(k, attr, value)

# Patch mockredis/mockredis/client.py
# The offical implementation assume decode_responses=False
# Here we detect the option first and only encode when decode_responses=False
def _encode(self, value):
"Return a bytestring representation of the value. Taken from redis-py connection.py"
if isinstance(value, bytes):
return value
elif isinstance(value, (int, long)):
value = str(value).encode('utf-8')
elif isinstance(value, float):
value = repr(value).encode('utf-8')
elif not isinstance(value, basestring):
value = str(value).encode('utf-8')
elif not self.decode_responses:
value = value.encode('utf-8', 'strict')
return value

# Patch mockredis/mockredis/client.py
# The official implementation will filter out keys with a slash '/'
# ref: https://github.com/locationlabs/mockredis/blob/master/mockredis/client.py
Expand All @@ -118,20 +145,12 @@ def keys(self, pattern='*'):
import fnmatch
import re

# making sure the pattern is unicode/str.
try:
pattern = pattern.decode('utf-8')
# This throws an AttributeError in python 3, or an
# UnicodeEncodeError in python 2
except (AttributeError, UnicodeEncodeError):
pass

# Make regex out of glob styled pattern.
regex = fnmatch.translate(pattern)
regex = re.compile(regex)

# Find every key that matches the pattern
return [key for key in self.redis.keys() if regex.match(key.decode('utf-8'))]
return [key for key in self.redis.keys() if regex.match(key)]


swsssdk.interface.DBInterface._subscribe_keyspace_notification = _subscribe_keyspace_notification
Expand Down

0 comments on commit 8079558

Please sign in to comment.