feat: initialize governed datasources module
This commit is contained in:
351
tests/test_lifecycle.py
Normal file
351
tests/test_lifecycle.py
Normal file
@@ -0,0 +1,351 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from govoplan_core.auth import ApiPrincipal
|
||||
from govoplan_core.core.access import PrincipalRef
|
||||
from govoplan_core.core.datasources import (
|
||||
CAPABILITY_DATASOURCE_ORIGINS,
|
||||
DatasourceAccessError,
|
||||
DatasourceField,
|
||||
DatasourceOrigin,
|
||||
DatasourceOriginReadRequest,
|
||||
DatasourceOriginReadResult,
|
||||
DatasourceReadRequest,
|
||||
DatasourceStageInput,
|
||||
)
|
||||
from govoplan_core.db.base import Base, utcnow
|
||||
from govoplan_datasources.backend.db.models import (
|
||||
DatasourceMaterializationRecord,
|
||||
DatasourceRecord,
|
||||
DatasourceStageRecord,
|
||||
)
|
||||
from govoplan_datasources.backend.service import (
|
||||
CATALOGUE_READ_SCOPE,
|
||||
SOURCE_WRITE_SCOPE,
|
||||
STAGE_WRITE_SCOPE,
|
||||
SqlDatasourceProvider,
|
||||
)
|
||||
|
||||
|
||||
def principal(
|
||||
tenant_id: str = "tenant-1",
|
||||
*,
|
||||
scopes: tuple[str, ...] = (
|
||||
CATALOGUE_READ_SCOPE,
|
||||
SOURCE_WRITE_SCOPE,
|
||||
STAGE_WRITE_SCOPE,
|
||||
),
|
||||
) -> ApiPrincipal:
|
||||
return ApiPrincipal(
|
||||
principal=PrincipalRef(
|
||||
account_id="account-1",
|
||||
membership_id="membership-1",
|
||||
tenant_id=tenant_id,
|
||||
scopes=frozenset(scopes),
|
||||
),
|
||||
account=object(),
|
||||
user=object(),
|
||||
)
|
||||
|
||||
|
||||
class FakeOriginProvider:
|
||||
def __init__(self) -> None:
|
||||
self.rows = [{"id": 1, "name": "Initial"}]
|
||||
|
||||
def _origin(self) -> DatasourceOrigin:
|
||||
return DatasourceOrigin(
|
||||
ref="snapshot:origin-1",
|
||||
source_name="connector_cases",
|
||||
name="Connector cases",
|
||||
kind="database",
|
||||
shape="tabular",
|
||||
supported_modes=("live", "cached"),
|
||||
provider="connectors.test",
|
||||
schema=(
|
||||
DatasourceField(name="id", data_type="integer", nullable=False),
|
||||
DatasourceField(name="name", data_type="string", nullable=False),
|
||||
),
|
||||
schema_version="1",
|
||||
fingerprint=f"version-{len(self.rows)}-{self.rows[-1]['name']}",
|
||||
row_count=len(self.rows),
|
||||
updated_at=utcnow(),
|
||||
)
|
||||
|
||||
def list_origins(
|
||||
self,
|
||||
_session: object,
|
||||
_principal: object,
|
||||
*,
|
||||
query: str = "",
|
||||
limit: int = 100,
|
||||
):
|
||||
origin = self._origin()
|
||||
return (origin,) if query.lower() in origin.name.lower() else ()
|
||||
|
||||
def get_origin(
|
||||
self,
|
||||
_session: object,
|
||||
_principal: object,
|
||||
*,
|
||||
origin_ref: str,
|
||||
):
|
||||
return self._origin() if origin_ref == "snapshot:origin-1" else None
|
||||
|
||||
def read_origin(
|
||||
self,
|
||||
_session: object,
|
||||
_principal: object,
|
||||
*,
|
||||
request: DatasourceOriginReadRequest,
|
||||
) -> DatasourceOriginReadResult:
|
||||
origin = self.get_origin(None, None, origin_ref=request.origin_ref)
|
||||
if origin is None:
|
||||
raise AssertionError("Unexpected origin")
|
||||
rows = self.rows[request.offset : request.offset + request.limit]
|
||||
return DatasourceOriginReadResult(
|
||||
origin=origin,
|
||||
rows=tuple(dict(row) for row in rows),
|
||||
total_rows=len(self.rows),
|
||||
truncated=request.offset + len(rows) < len(self.rows),
|
||||
)
|
||||
|
||||
|
||||
class FakeRegistry:
|
||||
def __init__(self, origin_provider: FakeOriginProvider) -> None:
|
||||
self.origin_provider = origin_provider
|
||||
|
||||
def has_capability(self, name: str) -> bool:
|
||||
return name == CAPABILITY_DATASOURCE_ORIGINS
|
||||
|
||||
def capability(self, name: str) -> object:
|
||||
if not self.has_capability(name):
|
||||
raise KeyError(name)
|
||||
return self.origin_provider
|
||||
|
||||
|
||||
class DatasourceLifecycleTests(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.engine = create_engine("sqlite:///:memory:")
|
||||
Base.metadata.create_all(
|
||||
self.engine,
|
||||
tables=[
|
||||
DatasourceRecord.__table__,
|
||||
DatasourceMaterializationRecord.__table__,
|
||||
DatasourceStageRecord.__table__,
|
||||
],
|
||||
)
|
||||
self.Session = sessionmaker(bind=self.engine)
|
||||
self.session = self.Session()
|
||||
self.origins = FakeOriginProvider()
|
||||
self.provider = SqlDatasourceProvider(
|
||||
registry=FakeRegistry(self.origins),
|
||||
)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self.session.close()
|
||||
Base.metadata.drop_all(
|
||||
self.engine,
|
||||
tables=[
|
||||
DatasourceStageRecord.__table__,
|
||||
DatasourceMaterializationRecord.__table__,
|
||||
DatasourceRecord.__table__,
|
||||
],
|
||||
)
|
||||
self.engine.dispose()
|
||||
|
||||
def test_static_stage_promote_update_and_frozen_read(self) -> None:
|
||||
first_stage = self.provider.create_stage(
|
||||
self.session,
|
||||
principal(),
|
||||
stage=DatasourceStageInput(
|
||||
name="Monthly cases",
|
||||
source_name="monthly_cases",
|
||||
kind="upload",
|
||||
mode="static",
|
||||
shape="tabular",
|
||||
rows=({"id": 1, "status": "new"},),
|
||||
),
|
||||
)
|
||||
datasource, first = self.provider.promote_stage(
|
||||
self.session,
|
||||
principal(),
|
||||
stage_ref=first_stage.ref,
|
||||
)
|
||||
frozen = self.provider.freeze_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
datasource_ref=datasource.ref,
|
||||
label="Import evidence",
|
||||
)
|
||||
second_stage = self.provider.create_stage(
|
||||
self.session,
|
||||
principal(),
|
||||
stage=DatasourceStageInput(
|
||||
name="Monthly cases",
|
||||
source_name="monthly_cases",
|
||||
kind="upload",
|
||||
mode="static",
|
||||
shape="tabular",
|
||||
target_datasource_ref=datasource.ref,
|
||||
rows=({"id": 2, "status": "checked", "note": "schema changed"},),
|
||||
),
|
||||
)
|
||||
updated, second = self.provider.promote_stage(
|
||||
self.session,
|
||||
principal(),
|
||||
stage_ref=second_stage.ref,
|
||||
)
|
||||
self.session.commit()
|
||||
|
||||
current = self.provider.read_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
request=DatasourceReadRequest(datasource_ref=datasource.ref),
|
||||
)
|
||||
frozen_result = self.provider.read_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
request=DatasourceReadRequest(
|
||||
datasource_ref=datasource.ref,
|
||||
consistency="frozen",
|
||||
),
|
||||
)
|
||||
|
||||
self.assertEqual(1, first.revision)
|
||||
self.assertEqual(2, frozen.revision)
|
||||
self.assertEqual(3, second.revision)
|
||||
self.assertEqual(
|
||||
[{"id": 2, "status": "checked", "note": "schema changed"}],
|
||||
list(current.rows),
|
||||
)
|
||||
self.assertEqual([{"id": 1, "status": "new"}], list(frozen_result.rows))
|
||||
self.assertEqual(second.ref, updated.current_materialization_ref)
|
||||
self.assertEqual("2", updated.schema_version)
|
||||
self.assertEqual(
|
||||
"2",
|
||||
current.datasource.schema_version,
|
||||
)
|
||||
self.assertEqual(frozen.ref, frozen_result.materialization.ref)
|
||||
|
||||
def test_live_reads_origin_and_cached_refresh_is_explicit(self) -> None:
|
||||
live = self.provider.register_origin(
|
||||
self.session,
|
||||
principal(),
|
||||
origin_ref="snapshot:origin-1",
|
||||
name="Live cases",
|
||||
source_name="live_cases",
|
||||
mode="live",
|
||||
)
|
||||
cached = self.provider.register_origin(
|
||||
self.session,
|
||||
principal(),
|
||||
origin_ref="snapshot:origin-1",
|
||||
name="Cached cases",
|
||||
source_name="cached_cases",
|
||||
mode="cached",
|
||||
)
|
||||
self.session.commit()
|
||||
|
||||
self.origins.rows = [
|
||||
{"id": 1, "name": "Initial"},
|
||||
{"id": 2, "name": "New"},
|
||||
]
|
||||
live_read = self.provider.read_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
request=DatasourceReadRequest(datasource_ref=live.ref),
|
||||
)
|
||||
cached_before = self.provider.read_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
request=DatasourceReadRequest(datasource_ref=cached.ref),
|
||||
)
|
||||
cached_live = self.provider.read_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
request=DatasourceReadRequest(
|
||||
datasource_ref=cached.ref,
|
||||
consistency="live",
|
||||
),
|
||||
)
|
||||
refreshed, materialization = self.provider.refresh_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
datasource_ref=cached.ref,
|
||||
)
|
||||
cached_after = self.provider.read_datasource(
|
||||
self.session,
|
||||
principal(),
|
||||
request=DatasourceReadRequest(datasource_ref=cached.ref),
|
||||
)
|
||||
|
||||
self.assertEqual(2, live_read.total_rows)
|
||||
self.assertEqual(1, cached_before.total_rows)
|
||||
self.assertEqual(2, cached_live.total_rows)
|
||||
self.assertEqual(2, cached_after.total_rows)
|
||||
self.assertEqual(materialization.ref, refreshed.current_materialization_ref)
|
||||
|
||||
def test_tenant_and_scope_isolation(self) -> None:
|
||||
stage = self.provider.create_stage(
|
||||
self.session,
|
||||
principal(),
|
||||
stage=DatasourceStageInput(
|
||||
name="Private",
|
||||
source_name="private",
|
||||
kind="upload",
|
||||
mode="static",
|
||||
shape="tabular",
|
||||
rows=({"id": 1},),
|
||||
),
|
||||
)
|
||||
datasource, _ = self.provider.promote_stage(
|
||||
self.session,
|
||||
principal(),
|
||||
stage_ref=stage.ref,
|
||||
)
|
||||
self.session.commit()
|
||||
|
||||
self.assertEqual(
|
||||
(),
|
||||
self.provider.list_datasources(self.session, principal("tenant-2")),
|
||||
)
|
||||
self.assertIsNone(
|
||||
self.provider.get_datasource(
|
||||
self.session,
|
||||
principal("tenant-2"),
|
||||
datasource_ref=datasource.ref,
|
||||
)
|
||||
)
|
||||
with self.assertRaises(DatasourceAccessError):
|
||||
self.provider.list_datasources(
|
||||
self.session,
|
||||
principal(scopes=()),
|
||||
)
|
||||
|
||||
def test_stage_writer_can_list_stages_without_catalogue_scope(self) -> None:
|
||||
writer = principal(scopes=(STAGE_WRITE_SCOPE,))
|
||||
stage = self.provider.create_stage(
|
||||
self.session,
|
||||
writer,
|
||||
stage=DatasourceStageInput(
|
||||
name="Pending import",
|
||||
source_name="pending_import",
|
||||
kind="upload",
|
||||
mode="static",
|
||||
shape="tabular",
|
||||
rows=({"id": 1},),
|
||||
),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
(stage.ref,),
|
||||
tuple(item.ref for item in self.provider.list_stages(self.session, writer)),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user