<gh_stars>0
#!/usr/bin/env python
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import os

from absl.testing import absltest
import mock

from grr_response_core import config
from grr_response_core.lib import flags
from grr_response_server import artifact_registry
from grr_response_server import data_store
from grr_response_server import db
from grr_response_server.databases import mem
from grr_response_server.databases import migration
from grr.test_lib import test_lib


def _SetUpArtifacts():
  test_artifacts_file = os.path.join(config.CONFIG["Test.data_dir"],
                                     "artifacts", "test_artifacts.json")
  artifact_registry.REGISTRY.AddFileSource(test_artifacts_file)


class ArtifactMigrationTest(absltest.TestCase):

  def setUp(self):
    super(ArtifactMigrationTest, self).setUp()

    self._db_patcher = mock.patch.object(
        data_store, "REL_DB", db.DatabaseValidationWrapper(mem.InMemoryDB()))
    self._db_patcher.start()

    self._artifact_patcher = mock.patch.object(
        artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry())
    self._artifact_patcher.start()

  def tearDown(self):
    self._db_patcher.stop()
    self._artifact_patcher.stop()
    super(ArtifactMigrationTest, self).tearDown()

  @mock.patch.object(data_store, "RelationalDBReadEnabled", return_value=False)
  def testMigratesAllArtifactsWithoutReadFromRelDB(self, unused_data_store):
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    _SetUpArtifacts()
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    migration.MigrateArtifacts()
    self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 29)

  @mock.patch.object(data_store, "RelationalDBReadEnabled", return_value=True)
  def testMigratesAllArtifactsWithReadFromRelDB(self, unused_data_store):
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    _SetUpArtifacts()
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    migration.MigrateArtifacts()
    self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 29)

  def testMigrationWithOverwriteDoesNotThrow(self):
    _SetUpArtifacts()
    migration.MigrateArtifacts()
    migration.MigrateArtifacts(overwrite=True)
    self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 29)


def main(args):
  test_lib.main(args)


if __name__ == "__main__":
  flags.StartMain(main)

#!/usr/bin/env python
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import os

from absl.testing import absltest
import mock

from grr_response_core import config
from grr_response_core.lib import flags
from grr_response_server import artifact_registry
from grr_response_server import data_store
from grr_response_server import db
from grr_response_server.databases import mem
from grr_response_server.databases import migration
from grr.test_lib import test_lib


def _SetUpArtifacts():
  test_artifacts_file = os.path.join(config.CONFIG["Test.data_dir"],
                                     "artifacts", "test_artifacts.json")
  artifact_registry.REGISTRY.AddFileSource(test_artifacts_file)


class ArtifactMigrationTest(absltest.TestCase):

  def setUp(self):
    super(ArtifactMigrationTest, self).setUp()

    self._db_patcher = mock.patch.object(
        data_store, "REL_DB", db.DatabaseValidationWrapper(mem.InMemoryDB()))
    self._db_patcher.start()

    self._artifact_patcher = mock.patch.object(
        artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry())
    self._artifact_patcher.start()

  def tearDown(self):
    self._db_patcher.stop()
    self._artifact_patcher.stop()
    super(ArtifactMigrationTest, self).tearDown()

  @mock.patch.object(data_store, "RelationalDBReadEnabled", return_value=False)
  def testMigratesAllArtifactsWithoutReadFromRelDB(self, unused_data_store):
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    _SetUpArtifacts()
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    migration.MigrateArtifacts()
    self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 29)

  @mock.patch.object(data_store, "RelationalDBReadEnabled", return_value=True)
  def testMigratesAllArtifactsWithReadFromRelDB(self, unused_data_store):
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    _SetUpArtifacts()
    self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
    migration.MigrateArtifacts()
    self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 29)

  def testMigrationWithOverwriteDoesNotThrow(self):
    _SetUpArtifacts()
    migration.MigrateArtifacts()
    migration.MigrateArtifacts(overwrite=True)
    self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 29)


def main(args):
  test_lib.main(args)


if __name__ == "__main__":
  flags.StartMain(main)
#!/usr/bin/env python