Szpitale-graph/szpitale-graph/src/test/java/com/developx/szpitale/load/GraphLoaderOrganizationIT.java

package com.developx.szpitale.load;

import com.developx.szpitale.model.*;
import com.developx.szpitale.model.enums.LinkBasis;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.neo4j.driver.AuthTokens;
import org.neo4j.driver.Driver;
import org.neo4j.driver.Session;
import org.testcontainers.containers.Neo4jContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;

import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.*;

@Testcontainers
class GraphLoaderOrganizationIT {

    @Container
    static Neo4jContainer<?> neo4j = new Neo4jContainer<>("neo4j:5")
            .withEnv("NEO4JLABS_PLUGINS", "['apoc']")
            .withEnv("NEO4J_apoc_export_file_enabled", "true")
            .withEnv("NEO4J_apoc_import_file_enabled", "true")
            .withEnv("NEO4J_apoc_import_file_use__neo4j__config", "true");

    static Driver driver;

    @BeforeAll
    static void setupDriver() {
        driver = org.neo4j.driver.GraphDatabase.driver(
                neo4j.getBoltUrl(),
                AuthTokens.basic("neo4j", "password"));
    }

    @Test
    void load_organizationAndCompany_nodesCreatedWithLinks() {
        GraphSchema schema = new GraphSchema(driver);
        schema.createConstraints();

        Organization org = new Organization(
            "org:test-fundacja",
            "Test Fundacja",
            "0000111222",
            null,
            List.of("https://example.org/fundacja")
        );

        Company company = new Company(
            "company:test-sp-zoo",
            "Test Sp. z o.o.",
            "0000333444",
            "1234567890",
            null,
            List.of("https://example.org/company")
        );

        Person person = new Person(
            "person:test-osoba",
            "Test Osoba",
            "Test Osoba",
            List.of(),
            List.of("https://example.org/person"),
            null, null
        );

        PersonOrganizationLink pol = new PersonOrganizationLink(
            "person:test-osoba",
            "org:test-fundacja",
            "Czonek Zarzadu",
            LinkBasis.KRS,
            List.of("https://example.org/link")
        );

        PersonCompanyLink pcl = new PersonCompanyLink(
            "person:test-osoba",
            "company:test-sp-zoo",
            "Wspolnik",
            LinkBasis.NIP,
            List.of("https://example.org/link2")
        );

        OrganizationLink orgLink = new OrganizationLink(
            "org:test-fundacja",
            "org:inna-fundacja",
            LinkBasis.NIP,
            "wspolny nip",
            List.of("https://example.org/link3")
        );

        CanonicalDataset dataset = new CanonicalDataset(
            "1.0", "test",
            java.time.LocalDateTime.now(),
            List.of(),
            List.of(person),
            List.of(), List.of(), List.of(),
            List.of(org),
            List.of(company),
            List.of(pol),
            List.of(orgLink),
            List.of(pcl)
        );

        GraphLoader loader = new GraphLoader(driver);
        GraphLoader.LoadStats stats = loader.load(List.of(dataset));

        assertEquals(1, stats.organizationCount);
        assertEquals(1, stats.companyCount);
        assertEquals(1, stats.personOrganizationLinkCount);
        assertEquals(1, stats.organizationLinkCount);
        assertEquals(1, stats.personCompanyLinkCount);

        try (Session session = driver.session()) {
            var orgResult = session.run(
                "MATCH (o:Organization {id: 'org:test-fundacja'}) RETURN o.name AS name").list();
            assertFalse(orgResult.isEmpty());
            assertEquals("Test Fundacja", orgResult.get(0).get("name").asString());

            var compResult = session.run(
                "MATCH (c:Company {id: 'company:test-sp-zoo'}) RETURN c.name AS name").list();
            assertFalse(compResult.isEmpty());
            assertEquals("Test Sp. z o.o.", compResult.get(0).get("name").asString());

            var polResult = session.run(
                "MATCH (p:Person {id: 'person:test-osoba'})-" +
                "[:CZLONEK_ORGANIZACJI]->(o:Organization) " +
                "RETURN o.id AS orgId").list();
            assertFalse(polResult.isEmpty());
            assertEquals("org:test-fundacja", polResult.get(0).get("orgId").asString());

            var pclResult = session.run(
                "MATCH (p:Person {id: 'person:test-osoba'})-" +
                "[:POWIAZANY_Z_FIRMA]->(c:Company) " +
                "RETURN c.id AS compId").list();
            assertFalse(pclResult.isEmpty());
            assertEquals("company:test-sp-zoo", pclResult.get(0).get("compId").asString());

            var orgLinkResult = session.run(
                "MATCH (a:Organization {id: 'org:test-fundacja'})-" +
                "[:POWIAZANA_Z]->(b:Organization) " +
                "RETURN b.id AS toId").list();
            assertFalse(orgLinkResult.isEmpty());
            assertEquals("org:inna-fundacja", orgLinkResult.get(0).get("toId").asString());
        }
    }

    @Test
    void schema_constraints_created() {
        GraphSchema schema = new GraphSchema(driver);
        schema.createConstraints();

        try (Session session = driver.session()) {
            var result = session.run("SHOW CONSTRAINTS");
            int count = 0;
            boolean hasOrgConstraint = false, hasCompConstraint = false;
            while (result.hasNext()) {
                count++;
                var record = result.next();
                String constraintName = record.get("name").asString();
                if (constraintName.contains("org_id")) {
                    hasOrgConstraint = true;
                }
                if (constraintName.contains("company_id")) {
                    hasCompConstraint = true;
                }
            }
            assertTrue(count > 0, "Should have constraints: found " + count);
            assertTrue(hasOrgConstraint, "Organization constraint should exist");
            assertTrue(hasCompConstraint, "Company constraint should exist");
        }
    }
}