text: # -*- coding: utf-8 -*-
  # ------------------------------------------------------------------------------
  #
  #   Copyright 2018-2019 Fetch.AI Limited
  #
  #   Licensed under the Apache License, Version 2.0 (the "License");
  #   you may not use this file except in compliance with the License.
  #   You may obtain a copy of the License at
  #
  #       http://www.apache.org/licenses/LICENSE-2.0
  #
  #   Unless required by applicable law or agreed to in writing, software
  #   distributed under the License is distributed on an "AS IS" BASIS,
  #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #   See the License for the specific language governing permissions and
  #   limitations under the License.
  #
  # ------------------------------------------------------------------------------
  """This module contains tests behaviour storage access."""
  import os
  from typing import List, Set
  
  import pytest
  
  from aea.aea import AEA
  from aea.aea_builder import AEABuilder
  from aea.configurations.base import SkillConfig
  from aea.configurations.constants import DEFAULT_LEDGER
  from aea.mail.base import Envelope
  from aea.protocols.base import Address, Message
  from aea.protocols.dialogue.base import (
      BasicDialoguesStorage,
      Dialogue,
      DialogueLabel,
      PersistDialoguesStorage,
  )
  from aea.skills.base import Handler, Skill, SkillContext
  from aea.skills.behaviours import TickerBehaviour
  from aea.test_tools.test_cases import AEATestCaseEmpty
  
  from packages.fetchai.protocols.default.dialogues import (
      DefaultDialogue,
      DefaultDialogues,
  )
  from packages.fetchai.protocols.default.message import DefaultMessage
  from packages.fetchai.skills.echo import PUBLIC_ID
  
  from tests.common.utils import wait_for_condition
  
  
  class TBehaviour(TickerBehaviour):
      """Simple behaviour to count how many acts were called."""
  
      OBJ_ID = "some"
      OBJ_BODY = {"data": 12}
      COL_NAME = "test"
  
      def setup(self) -> None:
          """Set up behaviour."""
          self.counter = 0
  
      def act(self) -> None:
          """Make an action."""
          if self.context.storage and self.context.storage.is_connected:
              col = self.context.storage.get_sync_collection(self.COL_NAME)
              col.put(self.OBJ_ID, self.OBJ_BODY)
          self.counter += 1
  
  
  class THandler(Handler):
      """Simple behaviour to count how many acts were called."""
  
      SUPPORTED_PROTOCOL = DefaultMessage.protocol_id
  
      OBJ_ID = "some"
      OBJ_BODY = {"data": 12}
      COL_NAME = "test"
  
      def setup(self) -> None:
          """Set up behaviour."""
          self.counter = 0
  
      def teardown(self) -> None:
          """Tear down handler."""
  
      def handle(self, *args, **kwargs) -> None:
          """Handle an evelope."""
          if self.context.storage and self.context.storage.is_connected:
              col = self.context.storage.get_sync_collection(self.COL_NAME)
              col.put(self.OBJ_ID, self.OBJ_BODY)
  
          self.counter += 1
  
  
  def test_storage_access_from_behaviour():
      """Test storage access from behaviour component."""
      builder = AEABuilder()
      builder.set_name("aea_1")
      builder.add_private_key(DEFAULT_LEDGER)
  
      skill_context = SkillContext()
      behaviour = TBehaviour(name="behaviour", skill_context=skill_context)
      test_skill = Skill(
          SkillConfig(name="test_skill", author="fetchai"),
          skill_context=skill_context,
          handlers={},
          behaviours={"behaviour": behaviour},
      )
  
      builder.add_component_instance(test_skill)
      builder.set_storage_uri("sqlite://:memory:")
      aea = builder.build()
      skill_context.set_agent_context(aea.context)
  
      aea.runtime._threaded = True
      aea.runtime.start()
  
      try:
          wait_for_condition(lambda: aea.is_running, timeout=10)
          wait_for_condition(lambda: behaviour.counter > 0, timeout=10)
  
          col = skill_context.storage.get_sync_collection(behaviour.COL_NAME)
          assert col.get(behaviour.OBJ_ID) == behaviour.OBJ_BODY
      finally:
          aea.runtime.stop()
          aea.runtime.wait_completed(sync=True, timeout=10)
  
  
  def test_storage_access_from_handler():
      """Test storage access from handler component."""
      builder = AEABuilder()
      builder.set_name("aea_1")
      builder.add_private_key(DEFAULT_LEDGER)
  
      skill_context = SkillContext()
      handler = THandler(name="behaviour", skill_context=skill_context)
      test_skill = Skill(
          SkillConfig(name="test_skill", author="fetchai"),
          skill_context=skill_context,
          handlers={"handler": handler},
          behaviours={},
      )
  
      builder.add_component_instance(test_skill)
      builder.set_storage_uri("sqlite://:memory:")
      aea = builder.build()
      skill_context.set_agent_context(aea.context)
  
      aea.runtime._threaded = True
      aea.runtime.start()
  
      msg = DefaultMessage(
          dialogue_reference=("", ""),
          message_id=1,
          target=0,
          performative=DefaultMessage.Performative.BYTES,
          content=b"hello",
      )
      msg.to = aea.identity.address
      msg.sender = aea.identity.address
      envelope = Envelope(
          to=msg.to, sender=msg.sender, protocol_id=msg.protocol_id, message=msg,
      )
      try:
          wait_for_condition(lambda: aea.is_running, timeout=10)
  
          aea.runtime.multiplexer.in_queue.put(envelope)
  
          wait_for_condition(lambda: handler.counter > 0, timeout=10)
  
          col = skill_context.storage.get_sync_collection(handler.COL_NAME)
          assert col.get(handler.OBJ_ID) == handler.OBJ_BODY
      finally:
          aea.runtime.stop()
          aea.runtime.wait_completed(sync=True, timeout=10)
  
  
  def _get_labels(dialogues: List[Dialogue]) -> Set[DialogueLabel]:
      return set([i.dialogue_label for i in dialogues])
  
  
  def _storage_all_dialogues_labels(storage: BasicDialoguesStorage) -> Set[DialogueLabel]:
      return _get_labels(
          storage.dialogues_in_active_state + storage.dialogues_in_terminal_state
      )
  
  
  class TestDialogueModelSaveLoad(AEATestCaseEmpty):
      """Test dialogues sved and loaded on agent restart."""
  
      def setup(self):
          """Set up the test case."""
          self.add_item("skill", "fetchai/echo:latest", local=True)
          pkey_file = os.path.join(self._get_cwd(), "privkey")
          self.generate_private_key("fetchai", pkey_file)
          self.add_private_key("fetchai", pkey_file, False)
          self.add_private_key("fetchai", pkey_file, True)
  
          def role_from_first_message(  # pylint: disable=unused-argument
              message: Message, receiver_address: Address
          ) -> Dialogue.Role:
              """Infer the role of the agent from an incoming/outgoing first message
  
              :param message: an incoming/outgoing first message
              :param receiver_address: the address of the receiving agent
              :return: The role of the agent
              """
              return DefaultDialogue.Role.AGENT
  
          self.dialogues = DefaultDialogues(
              self_address="another_agent",
              role_from_first_message=role_from_first_message,
          )
  
      def _build_aea(self) -> AEA:
          """Build an AEA."""
          builder = AEABuilder.from_aea_project(self._get_cwd())
          builder.set_storage_uri("sqlite://some_file.db")
          aea = builder.build()
          aea.runtime._threaded = True
          return aea
  
      def test_dialogues_dumped_and_restored_properly(self):
          """Test dialogues restored during restart of agent."""
          aea = self._build_aea()
          aea.runtime.start()
          try:
              wait_for_condition(lambda: aea.is_running, timeout=10)
              echo_skill = aea.resources.get_skill(PUBLIC_ID)
              assert (
                  not echo_skill.skill_context.default_dialogues._dialogues_storage._dialogues_by_dialogue_label
              )
              msg, dialogue = self.dialogues.create(
                  aea.name,
                  performative=DefaultMessage.Performative.BYTES,
                  content=b"hello",
              )
              envelope = Envelope(
                  to=msg.to, sender=msg.sender, protocol_id=msg.protocol_id, message=msg,
              )
              aea.runtime.multiplexer.in_queue.put(envelope)
  
              dialogue_storage: PersistDialoguesStorage = echo_skill.skill_context.default_dialogues._dialogues_storage
              wait_for_condition(
                  lambda: _storage_all_dialogues_labels(dialogue_storage), timeout=3,
              )
              dialogues_for_check = _storage_all_dialogues_labels(dialogue_storage)
          finally:
              aea.runtime.stop()
              aea.runtime.wait_completed(sync=True, timeout=10)
  
          aea = self._build_aea()
          aea.runtime.start()
          try:
              wait_for_condition(lambda: aea.is_running, timeout=10)
              echo_skill = aea.resources.get_skill(PUBLIC_ID)
  
              dialogue_storage: PersistDialoguesStorage = echo_skill.skill_context.default_dialogues._dialogues_storage
              wait_for_condition(
                  lambda: _storage_all_dialogues_labels(dialogue_storage), timeout=3,
              )
              assert (
                  _storage_all_dialogues_labels(dialogue_storage) == dialogues_for_check
              )
          finally:
              aea.runtime.stop()
              aea.runtime.wait_completed(sync=True, timeout=10)
  
  
  if __name__ == "__main__":
      pytest.main([os.path.basename(__file__)])
  