# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.


import pytest
import torch

from torchgeo.models import FarSeg


class TestFarSeg:
    @torch.no_grad()
    @pytest.mark.parametrize(
        "backbone,pretrained",
        [
            ("resnet18", True),
            ("resnet34", False),
            ("resnet50", True),
            ("resnet101", False),
        ],
    )
    def test_valid_backbone(self, backbone: str, pretrained: bool) -> None:
        model = FarSeg(classes=4, backbone=backbone, backbone_pretrained=pretrained)
        x = torch.randn(2, 3, 128, 128)
        y = model(x)

        assert y.shape == (2, 4, 128, 128)

    def test_invalid_backbone(self) -> None:
        match = "unknown backbone: anynet."
        with pytest.raises(ValueError, match=match):
            FarSeg(classes=4, backbone="anynet", backbone_pretrained=False)
