Skip to content

Commit

Permalink
Fix revert commands + websocket channel permissions (#678)
Browse files Browse the repository at this point in the history
  • Loading branch information
pl-buiquang authored Dec 15, 2021
1 parent b2cb61a commit 2f85ddc
Show file tree
Hide file tree
Showing 15 changed files with 161 additions and 34 deletions.
1 change: 0 additions & 1 deletion antarest/eventbus/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ def __init__(
self.start()

def push(self, event: Event) -> None:
# TODO add arg permissions with group/role, user, public
self.backend.push_event(event)

def add_listener(
Expand Down
5 changes: 2 additions & 3 deletions antarest/eventbus/web.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from antarest.core.config import Config
from antarest.core.interfaces.eventbus import IEventBus, Event
from antarest.core.jwt import JWTUser, DEFAULT_ADMIN_USER
from antarest.core.model import PermissionInfo, StudyPermissionType
from antarest.core.model import PermissionInfo, StudyPermissionType, PublicMode
from antarest.core.permissions import check_permission
from antarest.login.auth import Auth

Expand Down Expand Up @@ -82,7 +82,7 @@ async def broadcast(
self, message: str, permissions: PermissionInfo, channel: Optional[str]
) -> None:
for connection in self.active_connections:
if check_permission(
if channel is not None or check_permission(
connection.user, permissions, StudyPermissionType.READ
):
if (
Expand Down Expand Up @@ -119,7 +119,6 @@ async def connect(
raise HTTPException(status_code=HTTPStatus.UNAUTHORIZED)
user = Auth.get_user_from_token(token, jwt_manager)
if user is None:
# TODO check auth and subscribe to rooms
raise HTTPException(status_code=HTTPStatus.UNAUTHORIZED)
except Exception as e:
logger.error(
Expand Down
2 changes: 1 addition & 1 deletion antarest/matrixstore/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ def __eq__(self, other: Any) -> bool:
# https://github.com/samuelcolvin/pydantic/issues/1423
# https://github.com/samuelcolvin/pydantic/issues/1599
# https://github.com/samuelcolvin/pydantic/issues/1930
# TODO maybe we should reverting to only float because Any cause problem retrieving data from a node will have pandas forcing all to float anyway...
# Reverting to only float because Any cause problem retrieving data from a node will have pandas forcing all to float anyway...
# this cause matrix dump on disk (and then hash id) to be different for basically the same matrices
MatrixData = float

Expand Down
23 changes: 17 additions & 6 deletions antarest/study/storage/variantstudy/model/command/remove_area.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import logging
from typing import Any, List, Optional, Tuple, Dict

from antarest.core.model import JSON
Expand All @@ -6,6 +7,9 @@
FileStudyTreeConfig,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.model.command.common import (
CommandOutput,
CommandName,
Expand Down Expand Up @@ -199,12 +203,19 @@ def revert(
# todo revert binding constraints that has the area in constraint and also search in base for one
return [command]

area_commands, links_commands = (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_area(base, self.id)
return area_commands + links_commands
# todo revert binding constraints that has the area in constraint
try:
area_commands, links_commands = (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_area(base, self.id)
# todo revert binding constraints that has the area in constraint
return area_commands + links_commands
except ChildNotFoundError as e:
logging.getLogger(__name__).warning(
f"Failed to extract revert command for remove_area {self.id}",
exc_info=e,
)
return []

def _create_diff(self, other: "ICommand") -> List["ICommand"]:
return []
Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import logging
from typing import Any, List, Optional, Tuple, Dict

from antarest.core.model import JSON
Expand Down Expand Up @@ -98,10 +99,17 @@ def revert(
):
return [command]

return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_binding_constraint(base, self.id)
try:
return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_binding_constraint(base, self.id)
except Exception as e:
logging.getLogger(__name__).warning(
f"Failed to extract revert command for remove_binding_constraint {self.id}",
exc_info=e,
)
return []

def _create_diff(self, other: "ICommand") -> List["ICommand"]:
return []
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,14 @@
import logging
from typing import Any, List, Tuple, Dict

from antarest.study.storage.rawstudy.model.filesystem.config.model import (
transform_name_to_id,
FileStudyTreeConfig,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.model.command.common import (
CommandOutput,
CommandName,
Expand Down Expand Up @@ -175,11 +179,18 @@ def revert(
# todo revert binding constraints that has the cluster in constraint and also search in base for one
return [command]

return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_cluster(base, self.area_id, self.cluster_id)
# todo revert binding constraints that has the cluster in constraint
try:
return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_cluster(base, self.area_id, self.cluster_id)
# todo revert binding constraints that has the cluster in constraint
except ChildNotFoundError as e:
logging.getLogger(__name__).warning(
f"Failed to extract revert command for remove_cluster {self.area_id}#{self.cluster_id}",
exc_info=e,
)
return []

def _create_diff(self, other: "ICommand") -> List["ICommand"]:
return []
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,14 @@
import logging
from typing import Any, List, Optional, Tuple, Dict

from antarest.study.storage.rawstudy.model.filesystem.config.model import (
transform_name_to_id,
FileStudyTreeConfig,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.model.command.common import (
CommandOutput,
CommandName,
Expand Down Expand Up @@ -69,10 +73,17 @@ def revert(
and transform_name_to_id(command.name) == self.id
):
return [command]
return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_district(base, self.id)
try:
return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_district(base, self.id)
except Exception as e:
logging.getLogger(__name__).warning(
f"Failed to extract revert command for remove_district {self.id}",
exc_info=e,
)
return []

def _create_diff(self, other: "ICommand") -> List["ICommand"]:
return []
Expand Down
19 changes: 15 additions & 4 deletions antarest/study/storage/variantstudy/model/command/remove_link.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,13 @@
import logging
from typing import Any, List, Optional, Tuple, Dict

from antarest.study.storage.rawstudy.model.filesystem.config.model import (
FileStudyTreeConfig,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.model.command.common import (
CommandOutput,
CommandName,
Expand Down Expand Up @@ -115,10 +119,17 @@ def revert(
):
return [command]
area_from, area_to = sorted([self.area1, self.area2])
return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_link(base, area_from, area_to)
try:
return (
self.command_context.command_extractor
or CommandExtraction(self.command_context.matrix_service)
).extract_link(base, area_from, area_to)
except ChildNotFoundError as e:
logging.getLogger(__name__).warning(
f"Failed to extract revert command for remove_link {self.area1}/{self.area2}",
exc_info=e,
)
return []

def _create_diff(self, other: "ICommand") -> List["ICommand"]:
return []
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,12 +50,6 @@ def extract(self, study: FileStudy) -> List[CommandDTO]:
study_tree, ["layers", "layers"]
)
)
# todo create something out of variant manager commands to replace single rawnode files ?
# study_commands.append(
# self._generate_update_config(
# study_tree, ["settings", "comments"]
# )
# )
stopwatch.log_elapsed(
lambda x: logger.info(f"General command extraction done in {x}s")
)
Expand Down
53 changes: 53 additions & 0 deletions tests/eventbus/test_websocket_manager.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
import asyncio
from unittest import IsolatedAsyncioTestCase
from unittest.mock import Mock, MagicMock

from starlette.websockets import WebSocket

from antarest.core.jwt import JWTUser
from antarest.core.model import PermissionInfo, PublicMode
from antarest.eventbus.web import (
ConnectionManager,
WebsocketMessage,
WebsocketMessageAction,
)


class AsyncMock(MagicMock):
async def __call__(self, *args, **kwargs):
return super(AsyncMock, self).__call__(*args, **kwargs)


class ConnectionManagerTest(IsolatedAsyncioTestCase):
async def test_subscriptions(self):
ws_manager = ConnectionManager()

user = JWTUser(id=1, type="user", impersonator=1, groups=[])
subscription_message = WebsocketMessage(
action=WebsocketMessageAction.SUBSCRIBE, payload="foo"
)
unsubscription_message = WebsocketMessage(
action=WebsocketMessageAction.UNSUBSCRIBE, payload="foo"
)
mock_connection = AsyncMock(spec=WebSocket)
await ws_manager.connect(mock_connection, user)
assert len(ws_manager.active_connections) == 1

ws_manager.process_message(
subscription_message.json(), mock_connection, user
)
assert len(ws_manager.active_connections[0].channel_subscriptions) == 1
assert (
ws_manager.active_connections[0].channel_subscriptions[0] == "foo"
)

await ws_manager.broadcast("hello", PermissionInfo(), channel="foo")
mock_connection.send_text.assert_called_with("hello")

ws_manager.process_message(
unsubscription_message.json(), mock_connection, user
)
assert len(ws_manager.active_connections[0].channel_subscriptions) == 0

ws_manager.disconnect(mock_connection)
assert len(ws_manager.active_connections) == 0
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@
from antarest.matrixstore.service import MatrixService
from antarest.study.storage.rawstudy.io.reader import IniReader
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.business.matrix_constants_generator import (
GeneratorMatrixConstants,
)
Expand Down Expand Up @@ -382,6 +385,9 @@ def test_revert(command_context: CommandContext):
command_context=command_context,
)
]
base.command_context.command_extractor.extract_binding_constraint.side_effect = (
ChildNotFoundError()
)
base.revert([], study)
base.command_context.command_extractor.extract_binding_constraint.assert_called_with(
study, "foo"
Expand Down
6 changes: 6 additions & 0 deletions tests/variantstudy/model/command/test_manage_district.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
transform_name_to_id,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.business.matrix_constants_generator import (
GeneratorMatrixConstants,
)
Expand Down Expand Up @@ -193,6 +196,9 @@ def test_revert(command_context: CommandContext):

base = RemoveDistrict(id="id", command_context=command_context)
study = FileStudy(config=Mock(), tree=Mock())
base.command_context.command_extractor.extract_district.side_effect = (
ChildNotFoundError()
)
base.revert([], study)
base.command_context.command_extractor.extract_district.assert_called_with(
study, "id"
Expand Down
6 changes: 6 additions & 0 deletions tests/variantstudy/model/command/test_remove_area.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
transform_name_to_id,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.business.matrix_constants_generator import (
GeneratorMatrixConstants,
)
Expand Down Expand Up @@ -107,6 +110,9 @@ def test_revert(command_context: CommandContext):
[Mock()],
[Mock()],
)
base.command_context.command_extractor.extract_area.side_effect = (
ChildNotFoundError()
)
base.revert([], study)
base.command_context.command_extractor.extract_area.assert_called_with(
study, "foo"
Expand Down
6 changes: 6 additions & 0 deletions tests/variantstudy/model/command/test_remove_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
transform_name_to_id,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.business.matrix_constants_generator import (
GeneratorMatrixConstants,
)
Expand Down Expand Up @@ -135,6 +138,9 @@ def test_revert(command_context: CommandContext):
)
]
study = FileStudy(config=Mock(), tree=Mock())
base.command_context.command_extractor.extract_cluster.side_effect = (
ChildNotFoundError()
)
base.revert([], study)
base.command_context.command_extractor.extract_cluster.assert_called_with(
study, "foo", "bar"
Expand Down
6 changes: 6 additions & 0 deletions tests/variantstudy/model/command/test_remove_link.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
transform_name_to_id,
)
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.rawstudy.model.filesystem.folder_node import (
ChildNotFoundError,
)
from antarest.study.storage.variantstudy.business.default_values import (
FilteringOptions,
LinkProperties,
Expand Down Expand Up @@ -104,6 +107,9 @@ def test_revert(command_context: CommandContext):
area1="foo", area2="bar", command_context=command_context
)
study = FileStudy(config=Mock(), tree=Mock())
base.command_context.command_extractor.extract_link.side_effect = (
ChildNotFoundError()
)
base.revert([], study)
base.command_context.command_extractor.extract_link.assert_called_with(
study, "bar", "foo"
Expand Down

0 comments on commit 2f85ddc

Please sign in to comment.