diff --git a/src/server/tasks/knowledgegraph/api.py b/src/server/tasks/knowledgegraph/api.py index 48a5f9d..b9d3aa6 100644 --- a/src/server/tasks/knowledgegraph/api.py +++ b/src/server/tasks/knowledgegraph/api.py @@ -57,7 +57,7 @@ def get_relations(self, variable: Union[Variable, str]): Returns (None, "Observation: [...]"). """ if not isinstance(variable, Variable): - if not re.match(r'^([mf])\.[\w_]+$', variable): + if not re.match(r'^[mg]\.[\w_]+$', variable): raise ValueError("get_relations: variable must be a variable or an entity") cache_key = (self.task_id, variable if isinstance(variable, str) else hash(variable)) @@ -92,7 +92,7 @@ def get_neighbors(self, variable: Union[Variable, str], relation: str): Get neighbors via a relation. Returns (new_variable, "Observation: ..."). """ if not isinstance(variable, Variable): - if not re.match(r'^([mf])\.[\w_]+$', variable): + if not re.match(r'^[mg]\.[\w_]+$', variable): raise ValueError("get_neighbors: variable must be a variable or an entity") cache_key = (self.task_id, variable if isinstance(variable, str) else hash(variable)) diff --git a/tests/test_knowledgegraph_entity_ids.py b/tests/test_knowledgegraph_entity_ids.py new file mode 100644 index 0000000..2dfe66f --- /dev/null +++ b/tests/test_knowledgegraph_entity_ids.py @@ -0,0 +1,48 @@ +import unittest + +from src.server.tasks.knowledgegraph import api as knowledgegraph_api + + +class StubSparqlExecutor: + def __init__(self): + self.entities = [] + + def get_out_relations(self, entity): + self.entities.append(entity) + return [] + + +class KnowledgeGraphEntityIdTest(unittest.TestCase): + def setUp(self): + knowledgegraph_api.relation_cache.clear() + knowledgegraph_api.variable_relations_cache.clear() + self.executor = StubSparqlExecutor() + self.api = knowledgegraph_api.API(self.executor) + + def test_get_relations_accepts_g_prefixed_entity(self): + entity = "g.11b5lzm6b0" + + self.api.get_relations(entity) + + self.assertEqual(self.executor.entities, [entity]) + + def test_get_neighbors_accepts_g_prefixed_entity(self): + entity = "g.11b5lzm6b0" + relation = "test.relation" + knowledgegraph_api.variable_relations_cache[entity] = [relation] + knowledgegraph_api.range_info[relation] = "test.type" + self.addCleanup(knowledgegraph_api.range_info.pop, relation) + + variable, _ = self.api.get_neighbors(entity, relation) + + self.assertEqual(variable.program, f"(JOIN {relation}_inv {entity})") + + def test_f_prefixed_value_is_rejected(self): + with self.assertRaises(ValueError): + self.api.get_relations("f.not_an_entity") + with self.assertRaises(ValueError): + self.api.get_neighbors("f.not_an_entity", "test.relation") + + +if __name__ == "__main__": + unittest.main()