aboutsummaryrefslogtreecommitdiff
path: root/tests/mlia/test_core_data_analysis.py
blob: a782159ee2a79c4555ebd8d54d70d2a2024f08b6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Tests for module data_analysis."""
from dataclasses import dataclass

from mlia.core.common import DataItem
from mlia.core.data_analysis import Fact
from mlia.core.data_analysis import FactExtractor


def test_fact_extractor() -> None:
    """Test fact extractor."""

    @dataclass
    class SampleFact(Fact):
        """Sample fact."""

        msg: str

    class SampleExtractor(FactExtractor):
        """Sample extractor."""

        def analyze_data(self, data_item: DataItem) -> None:
            self.add_fact(SampleFact(f"Fact for {data_item}"))

    extractor = SampleExtractor()
    extractor.analyze_data(42)
    extractor.analyze_data("some data")

    facts = extractor.get_analyzed_data()
    assert facts == [SampleFact("Fact for 42"), SampleFact("Fact for some data")]