#!/usr/bin/env python
"""Tests for grr.lib.output_plugin."""
from unittest import mock

from absl import app

from grr_response_core.lib import rdfvalue
from grr_response_core.lib import registry
from grr_response_server import output_plugin
from grr_response_server.flows.general import transfer
from grr_response_server.rdfvalues import flow_runner as rdf_flow_runner
from grr_response_server.rdfvalues import output_plugin as rdf_output_plugin
from grr.test_lib import test_lib


class TestOutputPluginWithArgs(output_plugin.OutputPlugin):

  args_type = rdf_flow_runner.FlowRunnerArgs

  def ProcessResponses(self, state, responses):
    pass


class OutputPluginTest(test_lib.GRRBaseTest):

  def testGetArgsHandlesMissingPluginsCorrectly(self):
    plugin_args = rdf_flow_runner.FlowRunnerArgs(
        flow_name=transfer.GetFile.__name__)
    descriptor = rdf_output_plugin.OutputPluginDescriptor(
        plugin_name="TestOutputPluginWithArgs", args=plugin_args)
    serialized = descriptor.SerializeToBytes()

    deserialized = rdf_output_plugin.OutputPluginDescriptor.FromSerializedBytes(
        serialized)
    self.assertEqual(deserialized, descriptor)
    self.assertEqual(deserialized.GetPluginClass(), TestOutputPluginWithArgs)

    opr = registry.OutputPluginRegistry
    with mock.patch.object(opr, "PLUGIN_REGISTRY", opr.PLUGIN_REGISTRY.copy()):
      del opr.PLUGIN_REGISTRY["TestOutputPluginWithArgs"]

      deserialized = rdf_output_plugin.OutputPluginDescriptor.FromSerializedBytes(
          serialized)

      self.assertEqual(deserialized.GetPluginClass(),
                       output_plugin.UnknownOutputPlugin)
      # UnknownOutputPlugin should just return serialized arguments as bytes.
      self.assertEqual(deserialized.GetPluginArgsClass(), rdfvalue.RDFBytes)

      self.assertEqual(deserialized.args, descriptor.args.SerializeToBytes())


def main(argv):
  test_lib.main(argv)


if __name__ == "__main__":
  app.run(main)
