class ServiceStateWrapper: """ Class wraps ServiceState interactions with redis """ # Unique typename for Redis key REDIS_VALUE_TYPE = "systemd_status" def __init__(self): serde = RedisSerde(self.REDIS_VALUE_TYPE, get_proto_serializer(), get_proto_deserializer(ServiceExitStatus)) self._flat_dict = RedisFlatDict(get_default_client(), serde) def update_service_status(self, service_name: str, service_status: ServiceExitStatus) -> None: """ Update the service exit status for a given service """ if service_name in self._flat_dict: current_service_status = self._flat_dict[service_name] else: current_service_status = ServiceExitStatus() if service_status.latest_service_result == \ ServiceExitStatus.ServiceResult.Value("SUCCESS"): service_status.num_clean_exits = \ current_service_status.num_clean_exits + 1 service_status.num_fail_exits = \ current_service_status.num_fail_exits else: service_status.num_fail_exits = \ current_service_status.num_fail_exits + 1 service_status.num_clean_exits = \ current_service_status.num_clean_exits self._flat_dict[service_name] = service_status def get_service_status(self, service_name: str) -> ServiceExitStatus: """ Get the service status protobuf for a given service @returns ServiceStatus protobuf object """ return self._flat_dict[service_name] def get_all_services_status(self) -> [str, ServiceExitStatus]: """ Get a dict of service name to service status @return dict of service_name to service map """ service_status = {} for k, v in self._flat_dict.items(): service_status[k] = v return service_status def cleanup_service_status(self) -> None: """ Cleanup service status for all services in redis, mostly using for testing """ self._flat_dict.clear()
class StateReplicatorTests(TestCase): @mock.patch("redis.Redis", MockRedis) def setUp(self): self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) service = MagicMock() service.config = { # Replicate arbitrary orc8r protos 'state_protos': [{ 'proto_file': 'orc8r.protos.common_pb2', 'proto_msg': 'NetworkID', 'redis_key': NID_TYPE, 'state_scope': 'network' }, { 'proto_file': 'orc8r.protos.common_pb2', 'proto_msg': 'IDList', 'redis_key': IDList_TYPE, 'state_scope': 'gateway' }, { 'proto_file': 'orc8r.protos.service303_pb2', 'proto_msg': 'LogVerbosity', 'redis_key': LOG_TYPE, 'state_scope': 'gateway' }], 'json_state': [{ 'redis_key': FOO_TYPE, 'state_scope': 'network' }] } service.loop = self.loop # Bind the rpc server to a free port self._rpc_server = grpc.server( futures.ThreadPoolExecutor(max_workers=10)) port = self._rpc_server.add_insecure_port('0.0.0.0:0') # Add the servicer self._servicer = DummyStateServer() self._servicer.add_to_server(self._rpc_server) self._rpc_server.start() # Create a rpc stub self.channel = grpc.insecure_channel('0.0.0.0:{}'.format(port)) serde1 = RedisSerde(NID_TYPE, get_proto_serializer(), get_proto_deserializer(NetworkID)) serde2 = RedisSerde(IDList_TYPE, get_proto_serializer(), get_proto_deserializer(IDList)) serde3 = RedisSerde(LOG_TYPE, get_proto_serializer(), get_proto_deserializer(LogVerbosity)) serde4 = RedisSerde(FOO_TYPE, get_json_serializer(), get_json_deserializer()) self.nid_client = RedisFlatDict(get_default_client(), serde1) self.idlist_client = RedisFlatDict(get_default_client(), serde2) self.log_client = RedisFlatDict(get_default_client(), serde3) self.foo_client = RedisFlatDict(get_default_client(), serde4) # Set up and start state replicating loop grpc_client_manager = GRPCClientManager( service_name="state", service_stub=StateServiceStub, max_client_reuse=60, ) garbage_collector = GarbageCollector(service, grpc_client_manager) self.state_replicator = StateReplicator( service=service, garbage_collector=garbage_collector, grpc_client_manager=grpc_client_manager, ) self.state_replicator.start() @mock.patch("redis.Redis", MockRedis) def tearDown(self): self._rpc_server.stop(None) self.state_replicator.stop() self.loop.close() def convert_msg_to_state(self, redis_state, is_proto=True): if is_proto: json_converted_state = MessageToDict(redis_state) serialized_json_state = json.dumps(json_converted_state) else: serialized_json_state = jsonpickle.encode(redis_state) return serialized_json_state.encode("utf-8") @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) def test_collect_states_to_replicate(self): async def test(): # Ensure setup is initialized properly self.nid_client.clear() self.idlist_client.clear() self.log_client.clear() self.foo_client.clear() key = 'id1' self.nid_client[key] = NetworkID(id='foo') self.idlist_client[key] = IDList(ids=['bar', 'blah']) self.foo_client[key] = Foo("boo", 3) exp1 = self.convert_msg_to_state(self.nid_client[key]) exp2 = self.convert_msg_to_state(self.idlist_client[key]) exp3 = self.convert_msg_to_state(self.foo_client[key], False) req = await self.state_replicator._collect_states_to_replicate() self.assertEqual(3, len(req.states)) for state in req.states: if state.type == NID_TYPE: self.assertEqual('id1', state.deviceID) self.assertEqual(1, state.version) self.assertEqual(exp1, state.value) elif state.type == IDList_TYPE: self.assertEqual('aaa-bbb:id1', state.deviceID) self.assertEqual(1, state.version) self.assertEqual(exp2, state.value) elif state.type == FOO_TYPE: self.assertEqual('id1', state.deviceID) self.assertEqual(1, state.version) self.assertEqual(exp3, state.value) else: self.fail("Unknown state type %s" % state.type) # Cancel the replicator's loop so there are no other activities self.state_replicator._periodic_task.cancel() self.loop.run_until_complete(test()) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_replicate_states_success(self, get_rpc_mock): async def test(): get_rpc_mock.return_value = self.channel # Ensure setup is initialized properly self.nid_client.clear() self.idlist_client.clear() self.log_client.clear() self.foo_client.clear() key = 'id1' foo = Foo("boo", 4) self.nid_client[key] = NetworkID(id='foo') self.idlist_client[key] = IDList(ids=['bar', 'blah']) self.foo_client[key] = foo # Increment version self.idlist_client[key] = IDList(ids=['bar', 'blah']) req = await self.state_replicator._collect_states_to_replicate() self.assertEqual(3, len(req.states)) # Ensure in-memory map updates properly await self.state_replicator._send_to_state_service(req) self.assertEqual(3, len(self.state_replicator._state_versions)) mem_key1 = make_mem_key('id1', NID_TYPE) mem_key2 = make_mem_key('aaa-bbb:id1', IDList_TYPE) mem_key3 = make_mem_key('id1', FOO_TYPE) self.assertEqual(1, self.state_replicator._state_versions[mem_key1]) self.assertEqual(2, self.state_replicator._state_versions[mem_key2]) self.assertEqual(1, self.state_replicator._state_versions[mem_key3]) # Now add new state and update some existing state key2 = 'id2' self.nid_client[key2] = NetworkID(id='bar') self.idlist_client[key] = IDList(ids=['bar', 'foo']) req = await self.state_replicator._collect_states_to_replicate() self.assertEqual(2, len(req.states)) # Ensure in-memory map updates properly await self.state_replicator._send_to_state_service(req) self.assertEqual(4, len(self.state_replicator._state_versions)) mem_key4 = make_mem_key('id2', NID_TYPE) self.assertEqual(1, self.state_replicator._state_versions[mem_key1]) self.assertEqual(3, self.state_replicator._state_versions[mem_key2]) self.assertEqual(1, self.state_replicator._state_versions[mem_key3]) self.assertEqual(1, self.state_replicator._state_versions[mem_key4]) # Cancel the replicator's loop so there are no other activities self.state_replicator._periodic_task.cancel() self.loop.run_until_complete(test()) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_unreplicated_states(self, get_grpc_mock): async def test(): get_grpc_mock.return_value = self.channel # Add initial state to be replicated self.nid_client.clear() self.idlist_client.clear() self.log_client.clear() self.foo_client.clear() key = 'id1' key2 = 'id2' self.nid_client[key] = NetworkID(id='foo') self.idlist_client[key] = IDList(ids=['bar', 'blah']) # Increment version self.idlist_client[key] = IDList(ids=['bar', 'blah']) # Set state that will be 'unreplicated' self.log_client[key2] = LogVerbosity(verbosity=5) req = await self.state_replicator._collect_states_to_replicate() self.assertEqual(3, len(req.states)) # Ensure in-memory map updates properly for successful replications await self.state_replicator._send_to_state_service(req) self.assertEqual(2, len(self.state_replicator._state_versions)) mem_key1 = make_mem_key('id1', NID_TYPE) mem_key2 = make_mem_key('aaa-bbb:id1', IDList_TYPE) self.assertEqual(1, self.state_replicator._state_versions[mem_key1]) self.assertEqual(2, self.state_replicator._state_versions[mem_key2]) # Now run again, ensuring only the state the wasn't replicated # will be sent again req = await self.state_replicator._collect_states_to_replicate() self.assertEqual(1, len(req.states)) self.assertEqual('aaa-bbb:id2', req.states[0].deviceID) self.assertEqual(LOG_TYPE, req.states[0].type) # Cancel the replicator's loop so there are no other activities self.state_replicator._periodic_task.cancel() self.loop.run_until_complete(test()) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_resync_success(self, get_grpc_mock): async def test(): get_grpc_mock.return_value = self.channel self.nid_client.clear() self.idlist_client.clear() self.log_client.clear() self.foo_client.clear() key = 'id1' # Set state that will be 'unsynced' self.nid_client[key] = NetworkID(id='foo') self.idlist_client[key] = IDList(ids=['bar', 'blah']) # Increment state's version self.idlist_client[key] = IDList(ids=['bar', 'blah']) await self.state_replicator._resync() self.assertEqual(True, self.state_replicator._has_resync_completed) self.assertEqual(1, len(self.state_replicator._state_versions)) mem_key = make_mem_key('aaa-bbb:id1', IDList_TYPE) self.assertEqual(2, self.state_replicator._state_versions[mem_key]) # Cancel the replicator's loop so there are no other activities self.state_replicator._periodic_task.cancel() self.loop.run_until_complete(test()) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_resync_failure(self, get_grpc_mock): async def test(): get_grpc_mock.return_value = self.channel self.nid_client.clear() self.idlist_client.clear() self.log_client.clear() self.foo_client.clear() # Set state that will trigger the RpcError log_key = 'id1' self.log_client[log_key] = LogVerbosity(verbosity=5) try: await self.state_replicator._resync() except grpc.RpcError: pass self.assertEqual(False, self.state_replicator._has_resync_completed) self.assertEqual(0, len(self.state_replicator._state_versions)) # Cancel the replicator's loop so there are no other activities self.state_replicator._periodic_task.cancel() self.loop.run_until_complete(test()) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_deleted_replicated_state(self, get_grpc_mock): async def test(): get_grpc_mock.return_value = self.channel self.nid_client.clear() self.idlist_client.clear() self.log_client.clear() self.foo_client.clear() key = 'id1' self.nid_client[key] = NetworkID(id='foo') req = await self.state_replicator._collect_states_to_replicate() self.assertEqual(1, len(req.states)) # Ensure in-memory map updates properly await self.state_replicator._send_to_state_service(req) self.assertEqual(1, len(self.state_replicator._state_versions)) mem_key1 = make_mem_key('id1', NID_TYPE) self.assertEqual(1, self.state_replicator._state_versions[mem_key1]) # Now delete state and ensure in-memory map gets updated properly del self.nid_client[key] req = await self.state_replicator._collect_states_to_replicate() self.assertIsNone(req) await self.state_replicator._cleanup_deleted_keys() self.assertFalse(key in self.state_replicator._state_versions) # Cancel the replicator's loop so there are no other activities self.state_replicator._periodic_task.cancel() self.loop.run_until_complete(test())
class RedisDictTests(TestCase): """ Tests for the RedisHashDict and RedisFlatDict containers """ @mock.patch("redis.Redis", MockRedis) def setUp(self): client = get_default_client() # Use arbitrary orc8r proto to test with self._hash_dict = RedisHashDict(client, "unittest", get_proto_serializer(), get_proto_deserializer(LogVerbosity)) serde = RedisSerde('log_verbosity', get_proto_serializer(), get_proto_deserializer(LogVerbosity)) self._flat_dict = RedisFlatDict(client, serde) @mock.patch("redis.Redis", MockRedis) def test_hash_insert(self): expected = LogVerbosity(verbosity=0) expected2 = LogVerbosity(verbosity=1) # insert proto self._hash_dict['key1'] = expected version = self._hash_dict.get_version("key1") actual = self._hash_dict['key1'] self.assertEqual(1, version) self.assertEqual(expected, actual) # update proto self._hash_dict['key1'] = expected2 version2 = self._hash_dict.get_version("key1") actual2 = self._hash_dict['key1'] self.assertEqual(2, version2) self.assertEqual(expected2, actual2) @mock.patch("redis.Redis", MockRedis) def test_missing_version(self): missing_version = self._hash_dict.get_version("key2") self.assertEqual(0, missing_version) @mock.patch("redis.Redis", MockRedis) def test_hash_delete(self): expected = LogVerbosity(verbosity=2) self._hash_dict['key3'] = expected actual = self._hash_dict['key3'] self.assertEqual(expected, actual) self._hash_dict.pop('key3') self.assertRaises(KeyError, self._hash_dict.__getitem__, 'key3') @mock.patch("redis.Redis", MockRedis) def test_flat_insert(self): expected = LogVerbosity(verbosity=5) expected2 = LogVerbosity(verbosity=1) # insert proto self._flat_dict['key1'] = expected version = self._flat_dict.get_version("key1") actual = self._flat_dict['key1'] self.assertEqual(1, version) self.assertEqual(expected, actual) # update proto self._flat_dict["key1"] = expected2 version2 = self._flat_dict.get_version("key1") actual2 = self._flat_dict["key1"] actual3 = self._flat_dict.get("key1") self.assertEqual(2, version2) self.assertEqual(expected2, actual2) self.assertEqual(expected2, actual3) @mock.patch("redis.Redis", MockRedis) def test_flat_missing_version(self): missing_version = self._flat_dict.get_version("key2") self.assertEqual(0, missing_version) @mock.patch("redis.Redis", MockRedis) def test_flat_bad_key(self): expected = LogVerbosity(verbosity=2) self.assertRaises(ValueError, self._flat_dict.__setitem__, 'bad:key', expected) self.assertRaises(ValueError, self._flat_dict.__getitem__, 'bad:key') self.assertRaises(ValueError, self._flat_dict.__delitem__, 'bad:key') @mock.patch("redis.Redis", MockRedis) def test_flat_delete(self): expected = LogVerbosity(verbosity=2) self._flat_dict['key3'] = expected actual = self._flat_dict['key3'] self.assertEqual(expected, actual) del self._flat_dict['key3'] self.assertRaises(KeyError, self._flat_dict.__getitem__, 'key3') self.assertEqual(None, self._flat_dict.get('key3')) @mock.patch("redis.Redis", MockRedis) def test_flat_clear(self): expected = LogVerbosity(verbosity=2) self._flat_dict['key3'] = expected actual = self._flat_dict['key3'] self.assertEqual(expected, actual) self._flat_dict.clear() self.assertEqual(0, len(self._flat_dict.keys())) @mock.patch("redis.Redis", MockRedis) def test_flat_garbage_methods(self): expected = LogVerbosity(verbosity=2) expected2 = LogVerbosity(verbosity=3) key = "k1" key2 = "k2" bad_key = "bad_key" self._flat_dict[key] = expected self._flat_dict[key2] = expected2 self._flat_dict.mark_as_garbage(key) is_garbage = self._flat_dict.is_garbage(key) self.assertTrue(is_garbage) is_garbage2 = self._flat_dict.is_garbage(key2) self.assertFalse(is_garbage2) self.assertEqual([key], self._flat_dict.garbage_keys()) self.assertEqual([key2], self._flat_dict.keys()) self.assertIsNone(self._flat_dict.get(key)) self.assertEqual(expected2, self._flat_dict.get(key2)) deleted = self._flat_dict.delete_garbage(key) not_deleted = self._flat_dict.delete_garbage(key2) self.assertTrue(deleted) self.assertFalse(not_deleted) self.assertIsNone(self._flat_dict.get(key)) self.assertEqual(expected2, self._flat_dict.get(key2)) with self.assertRaises(KeyError): self._flat_dict.is_garbage(bad_key) with self.assertRaises(KeyError): self._flat_dict.mark_as_garbage(bad_key)
class GarbageCollectorTests(TestCase): def setUp(self): self.mock_redis = fakeredis.FakeStrictRedis() self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) service = MagicMock() service.config = { # Replicate arbitrary orc8r protos 'state_protos': [{'proto_file': 'orc8r.protos.common_pb2', 'proto_msg': 'NetworkID', 'redis_key': NID_TYPE, 'state_scope': 'network'}, {'proto_file': 'orc8r.protos.service303_pb2', 'proto_msg': 'LogVerbosity', 'redis_key': LOG_TYPE, 'state_scope': 'gateway'}, ], 'json_state': [{'redis_key': FOO_TYPE, 'state_scope': 'gateway'}] } service.loop = self.loop # Bind the rpc server to a free port self._rpc_server = grpc.server( futures.ThreadPoolExecutor(max_workers=10) ) port = self._rpc_server.add_insecure_port('0.0.0.0:0') # Add the servicer self._servicer = DummyStateServer() self._servicer.add_to_server(self._rpc_server) self._rpc_server.start() # Create a rpc stub self.channel = grpc.insecure_channel('0.0.0.0:{}'.format(port)) serde1 = RedisSerde(NID_TYPE, get_proto_serializer(), get_proto_deserializer(NetworkID)) serde2 = RedisSerde(FOO_TYPE, get_json_serializer(), get_json_deserializer()) serde3 = RedisSerde(LOG_TYPE, get_proto_serializer(), get_proto_deserializer(LogVerbosity)) self.nid_client = RedisFlatDict(self.mock_redis, serde1) self.foo_client = RedisFlatDict(self.mock_redis, serde2) self.log_client = RedisFlatDict(self.mock_redis, serde3) # Set up and start garbage collecting loop grpc_client_manager = GRPCClientManager( service_name="state", service_stub=StateServiceStub, max_client_reuse=60, ) # mock the get_default_client function used to return the same # fakeredis object func_mock = mock.MagicMock(return_value=self.mock_redis) with patch('magma.state.redis_dicts.get_default_client', func_mock): # Start state garbage collection loop self.garbage_collector = GarbageCollector(service, grpc_client_manager) def tearDown(self): self._rpc_server.stop(None) self.loop.close() @mock.patch('snowflake.snowflake', get_mock_snowflake) def test_collect_states_to_delete(self): async def test(): # Ensure setup is initialized properly self.nid_client.clear() self.foo_client.clear() self.log_client.clear() key = 'id1' self.nid_client[key] = NetworkID(id='foo') self.foo_client[key] = Foo("boo", 3) req = await self.garbage_collector._collect_states_to_delete() self.assertIsNone(req) self.nid_client.mark_as_garbage(key) self.foo_client.mark_as_garbage(key) req = await self.garbage_collector._collect_states_to_delete() self.assertEqual(2, len(req.ids)) for state_id in req.ids: if state_id.type == NID_TYPE: self.assertEqual('id1', state_id.deviceID) elif state_id.type == FOO_TYPE: self.assertEqual('aaa-bbb:id1', state_id.deviceID) else: self.fail("Unknown state type %s" % state_id.type) # Cleanup del self.foo_client[key] del self.nid_client[key] self.loop.run_until_complete(test()) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_garbage_collect_success(self, get_rpc_mock): async def test(): get_rpc_mock.return_value = self.channel self.nid_client.clear() self.foo_client.clear() self.log_client.clear() key = 'id1' foo = Foo("boo", 4) self.nid_client[key] = NetworkID(id='foo') self.foo_client[key] = foo self.nid_client.mark_as_garbage(key) self.foo_client.mark_as_garbage(key) req = await self.garbage_collector._collect_states_to_delete() self.assertEqual(2, len(req.ids)) # Ensure all garbage collected objects get deleted from Redis await self.garbage_collector._send_to_state_service(req) self.assertEqual(0, len(self.nid_client.keys())) self.assertEqual(0, len(self.foo_client.keys())) self.assertEqual(0, len(self.nid_client.garbage_keys())) self.assertEqual(0, len(self.foo_client.garbage_keys())) self.loop.run_until_complete(test()) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_garbage_collect_rpc_failure(self, get_rpc_mock): async def test(): get_rpc_mock.return_value = self.channel self.nid_client.clear() self.foo_client.clear() self.log_client.clear() key = 'id1' self.nid_client[key] = NetworkID(id='foo') self.log_client[key] = LogVerbosity(verbosity=3) self.nid_client.mark_as_garbage(key) self.log_client.mark_as_garbage(key) req = await self.garbage_collector._collect_states_to_delete() self.assertEqual(2, len(req.ids)) # Ensure objects on deleted from Redis on RPC failure await self.garbage_collector._send_to_state_service(req) self.assertEqual(0, len(self.nid_client.keys())) self.assertEqual(0, len(self.log_client.keys())) self.assertEqual(1, len(self.nid_client.garbage_keys())) self.assertEqual(1, len(self.log_client.garbage_keys())) # Cleanup del self.log_client[key] del self.nid_client[key] self.loop.run_until_complete(test()) @mock.patch('snowflake.snowflake', get_mock_snowflake) @mock.patch('magma.magmad.state_reporter.ServiceRegistry.get_rpc_channel') def test_garbage_collect_with_state_update(self, get_rpc_mock): async def test(): get_rpc_mock.return_value = self.channel self.nid_client.clear() self.foo_client.clear() self.log_client.clear() key = 'id1' foo = Foo("boo", 4) self.nid_client[key] = NetworkID(id='foo') self.foo_client[key] = foo self.nid_client.mark_as_garbage(key) self.foo_client.mark_as_garbage(key) req = await self.garbage_collector._collect_states_to_delete() self.assertEqual(2, len(req.ids)) # Update one of the states, to ensure we don't delete valid state # from Redis expected = NetworkID(id='bar') self.nid_client[key] = expected # Ensure all garbage collected objects get deleted from Redis await self.garbage_collector._send_to_state_service(req) self.assertEqual(1, len(self.nid_client.keys())) self.assertEqual(0, len(self.foo_client.keys())) self.assertEqual(0, len(self.nid_client.garbage_keys())) self.assertEqual(0, len(self.foo_client.garbage_keys())) self.assertEqual(expected, self.nid_client[key]) self.loop.run_until_complete(test())