# Path: mass_flask_core/models/sample.py
# class Sample(db.Document, CommentsMixin):
#     delivery_date = DateTimeField(default=TimeFunctions.get_timestamp, required=True)
#     first_seen = DateTimeField(default=TimeFunctions.get_timestamp, required=True)
#     tags = ListField(StringField(regex=r'^[\w:\-\_\/\+\.]+$'))
#     dispatched_to = ListField(ReferenceField(AnalysisSystem))
#     created_by = GenericReferenceField()
#     tlp_level = TLPLevelField(verbose_name='TLP Level', help_text='Privacy level of this sample', required=True)
# 
#     meta = {
#         'allow_inheritance': True,
#         'ordering': ['-delivery_date'],
#         'indexes': ['delivery_date'],
#         'queryset_class': SampleQuerySet
#     }
# 
#     def __repr__(self):
#         return '[{}] {}'.format(str(self.__class__.__name__), str(self.id))
# 
#     def __str__(self):
#         return self.__repr__()
# 
#     @property
#     def title(self):
#         return self.id
# 
#     def add_tags(self, tags):
#         self.tags = ListFunctions.merge_lists_without_duplicates(self.tags, tags)
#         self.save()
# 
#     def _update(self, **kwargs):
#         if 'first_seen' in kwargs and kwargs['first_seen'] < self.first_seen:
#             self.first_seen = kwargs['first_seen']
#         if 'tlp_level' in kwargs and not self.tlp_level:
#             self.tlp_level = kwargs['tlp_level']
#         if current_authenticated_entity.is_authenticated and not self.created_by:
#             self.created_by = current_authenticated_entity._get_current_object()
#         if 'tags' in kwargs:
#             tags = kwargs['tags']
#         else:
#             tags = []
#         self.tags = ListFunctions.merge_lists_without_duplicates(self.tags, tags)
# 
#     @classmethod
#     def create_or_update(cls, **kwargs):
#         raise ValidationError('Samples of the base type Sample should not be instantiated directly.')
# 
# Path: mass_flask_core/models/sample_relation.py
# class SampleRelation(db.Document):
#     sample = ReferenceField(Sample, required=True)
#     other = ReferenceField(Sample, required=True)
#     meta = {
#         'allow_inheritance': True,
#     }
# 
#     def __repr__(self):
#         repr_string = '[{}] {} -- {}'
#         return repr_string.format(self.__class__.__name__, self.sample, self.other)
# 
#     def __str__(self):
#         return self.__repr__()
# 
#     def _initialize(self, **kwargs):
#         self.sample = Sample.objects(id=kwargs['sample']).first()
#         self.other = Sample.objects(id=kwargs['other']).first()
# 
#     @property
#     def title(self):
#         return self.id
# 
#     @classmethod
#     def create(cls, **kwargs):
#         if 'sample' not in kwargs or 'other' not in kwargs:
#             raise ValidationError('Parameter "sample" or "other" missing')
#         else:
#             sample_relation = cls()
#             sample_relation._initialize(**kwargs)
#             sample_relation.save()
#             return sample_relation
# 
# Path: mass_flask_core/tests/flask_test_case.py
# class FlaskTestCase(TestCase):
#     @staticmethod
#     def _clean_test_database():
#         conn = connection.get_connection()
#         db = connection.get_db()
#         conn.drop_database(db)
# 
#     def __call__(self, result=None):
#         self._pre_setup()
#         super(FlaskTestCase, self).__call__(result)
#         self._post_tearDown()
# 
#     def _pre_setup(self):
#         bootstrap_mass_flask()
#         if not 'MASS_TESTING' in app.config or app.config['MASS_TESTING'] is False:
#             raise RuntimeError('Running unit test without TESTING=True in configuration. Aborting.')
#         self.app = app
#         self.client = app.test_client()
#         self._ctx = self.app.test_request_context()
#         self._ctx.push()
# 
#     def _post_tearDown(self):
#         self._ctx.pop()
#         self._clean_test_database()
# 
# Path: mass_flask_core/utils/graph_functions.py
# class GraphFunctions:
#     @staticmethod
#     def get_relation_graph(sample, depth=3):
#         current_depth = 0
#         samples_at_current_depth = [sample.id]
#         all_queried_samples = []
#         found_relations = set()
# 
#         while current_depth < depth:
#             from mass_flask_core.models import SampleRelation
#             relations_at_current_depth = SampleRelation.objects(Q(sample__in=samples_at_current_depth) | Q(other__in=samples_at_current_depth)).no_dereference()
#             all_queried_samples.extend(samples_at_current_depth)
#             samples_at_next_depth = []
#             for relation in relations_at_current_depth:
#                 found_relations.add(relation)
#                 if relation.sample.id not in all_queried_samples:
#                     samples_at_next_depth.append(relation.sample.id)
#                 if relation.other.id not in all_queried_samples:
#                     samples_at_next_depth.append(relation.other.id)
#             samples_at_current_depth = samples_at_next_depth
#             current_depth += 1
#         return found_relations


from mixer.backend.mongoengine import mixer
from mass_flask_core.models import Sample, SampleRelation
from mass_flask_core.tests import FlaskTestCase
from mass_flask_core.utils import GraphFunctions



class RelationGraphTestCase(FlaskTestCase):
    def assertGraphEqual(self, sample, result_set, depth=None):
        if depth is None:
            self.assertEqual(GraphFunctions.get_relation_graph(sample), result_set)
        else:
            self.assertEqual(GraphFunctions.get_relation_graph(sample, depth=depth), result_set)

    def test_simple_graph(self):

# Path: mass_flask_core/models/sample.py
# class Sample(db.Document, CommentsMixin):
#     delivery_date = DateTimeField(default=TimeFunctions.get_timestamp, required=True)
#     first_seen = DateTimeField(default=TimeFunctions.get_timestamp, required=True)
#     tags = ListField(StringField(regex=r'^[\w:\-\_\/\+\.]+$'))
#     dispatched_to = ListField(ReferenceField(AnalysisSystem))
#     created_by = GenericReferenceField()
#     tlp_level = TLPLevelField(verbose_name='TLP Level', help_text='Privacy level of this sample', required=True)
# 
#     meta = {
#         'allow_inheritance': True,
#         'ordering': ['-delivery_date'],
#         'indexes': ['delivery_date'],
#         'queryset_class': SampleQuerySet
#     }
# 
#     def __repr__(self):
#         return '[{}] {}'.format(str(self.__class__.__name__), str(self.id))
# 
#     def __str__(self):
#         return self.__repr__()
# 
#     @property
#     def title(self):
#         return self.id
# 
#     def add_tags(self, tags):
#         self.tags = ListFunctions.merge_lists_without_duplicates(self.tags, tags)
#         self.save()
# 
#     def _update(self, **kwargs):
#         if 'first_seen' in kwargs and kwargs['first_seen'] < self.first_seen:
#             self.first_seen = kwargs['first_seen']
#         if 'tlp_level' in kwargs and not self.tlp_level:
#             self.tlp_level = kwargs['tlp_level']
#         if current_authenticated_entity.is_authenticated and not self.created_by:
#             self.created_by = current_authenticated_entity._get_current_object()
#         if 'tags' in kwargs:
#             tags = kwargs['tags']
#         else:
#             tags = []
#         self.tags = ListFunctions.merge_lists_without_duplicates(self.tags, tags)
# 
#     @classmethod
#     def create_or_update(cls, **kwargs):
#         raise ValidationError('Samples of the base type Sample should not be instantiated directly.')
# 
# Path: mass_flask_core/models/sample_relation.py
# class SampleRelation(db.Document):
#     sample = ReferenceField(Sample, required=True)
#     other = ReferenceField(Sample, required=True)
#     meta = {
#         'allow_inheritance': True,
#     }
# 
#     def __repr__(self):
#         repr_string = '[{}] {} -- {}'
#         return repr_string.format(self.__class__.__name__, self.sample, self.other)
# 
#     def __str__(self):
#         return self.__repr__()
# 
#     def _initialize(self, **kwargs):
#         self.sample = Sample.objects(id=kwargs['sample']).first()
#         self.other = Sample.objects(id=kwargs['other']).first()
# 
#     @property
#     def title(self):
#         return self.id
# 
#     @classmethod
#     def create(cls, **kwargs):
#         if 'sample' not in kwargs or 'other' not in kwargs:
#             raise ValidationError('Parameter "sample" or "other" missing')
#         else:
#             sample_relation = cls()
#             sample_relation._initialize(**kwargs)
#             sample_relation.save()
#             return sample_relation
# 
# Path: mass_flask_core/tests/flask_test_case.py
# class FlaskTestCase(TestCase):
#     @staticmethod
#     def _clean_test_database():
#         conn = connection.get_connection()
#         db = connection.get_db()
#         conn.drop_database(db)
# 
#     def __call__(self, result=None):
#         self._pre_setup()
#         super(FlaskTestCase, self).__call__(result)
#         self._post_tearDown()
# 
#     def _pre_setup(self):
#         bootstrap_mass_flask()
#         if not 'MASS_TESTING' in app.config or app.config['MASS_TESTING'] is False:
#             raise RuntimeError('Running unit test without TESTING=True in configuration. Aborting.')
#         self.app = app
#         self.client = app.test_client()
#         self._ctx = self.app.test_request_context()
#         self._ctx.push()
# 
#     def _post_tearDown(self):
#         self._ctx.pop()
#         self._clean_test_database()
# 
# Path: mass_flask_core/utils/graph_functions.py
# class GraphFunctions:
#     @staticmethod
#     def get_relation_graph(sample, depth=3):
#         current_depth = 0
#         samples_at_current_depth = [sample.id]
#         all_queried_samples = []
#         found_relations = set()
# 
#         while current_depth < depth:
#             from mass_flask_core.models import SampleRelation
#             relations_at_current_depth = SampleRelation.objects(Q(sample__in=samples_at_current_depth) | Q(other__in=samples_at_current_depth)).no_dereference()
#             all_queried_samples.extend(samples_at_current_depth)
#             samples_at_next_depth = []
#             for relation in relations_at_current_depth:
#                 found_relations.add(relation)
#                 if relation.sample.id not in all_queried_samples:
#                     samples_at_next_depth.append(relation.sample.id)
#                 if relation.other.id not in all_queried_samples:
#                     samples_at_next_depth.append(relation.other.id)
#             samples_at_current_depth = samples_at_next_depth
#             current_depth += 1
#         return found_relations


# Path: mass_flask_core/tests/test_relation_graph.py
from mixer.backend.mongoengine import mixer
from mass_flask_core.models import Sample, SampleRelation
from mass_flask_core.tests import FlaskTestCase
from mass_flask_core.utils import GraphFunctions




class RelationGraphTestCase(FlaskTestCase):
    def assertGraphEqual(self, sample, result_set, depth=None):
        if depth is None:
            self.assertEqual(GraphFunctions.get_relation_graph(sample), result_set)
        else:
            self.assertEqual(GraphFunctions.get_relation_graph(sample, depth=depth), result_set)

    def test_simple_graph(self):

        sample1 = mixer.blend(Sample)