Repository: telstra/open-kilda
Issue: Can't update one switch Y flow

**Steps to reproduce:**
1. create one switch Y flow
```
{
  "y_flow_id": "flow1",
  "shared_endpoint": {
    "port_number": 10,
    "switch_id": "2"
  },
  "maximum_bandwidth": 10000,
  "ignore_bandwidth": true,
  "sub_flows": [
    {
      "flow_id": "sub1",
      "endpoint": {
        "inner_vlan_id": 0,
        "port_number": 10,
        "switch_id": "2",
        "vlan_id": 123
      },
      "shared_endpoint": {
        "inner_vlan_id": 0,
        "vlan_id": 456
      }
    },
{
      "flow_id": "sub2",
      "endpoint": {
        "inner_vlan_id": 0,
        "port_number": 10,
        "switch_id": "2",
        "vlan_id": 345
      },
      "shared_endpoint": {
        "inner_vlan_id": 0,
        "vlan_id": 789
      }
    }
  ]
}
```
2. update flow bandwidth via PATCH request
```
{
  "maximum_bandwidth": 20000
}
```

**Expected result:**
flow was updated

**Actual result:**
Flow is stuck in "In progress" state

Changed files:
src-java/kilda-model/src/main/java/org/openkilda/model/YSubFlow.java
src-java/flowhs-topology/flowhs-storm-topology/src/main/java/org/openkilda/wfm/topology/flowhs/fsm/yflow/update/actions/ValidateYFlowAction.java

diff --git a/src-java/flowhs-topology/flowhs-storm-topology/src/main/java/org/openkilda/wfm/topology/flowhs/fsm/yflow/update/actions/ValidateYFlowAction.java b/src-java/flowhs-topology/flowhs-storm-topology/src/main/java/org/openkilda/wfm/topology/flowhs/fsm/yflow/update/actions/ValidateYFlowAction.java
index 9e2eb265f..89b78109f 100644
--- a/src-java/flowhs-topology/flowhs-storm-topology/src/main/java/org/openkilda/wfm/topology/flowhs/fsm/yflow/update/actions/ValidateYFlowAction.java
+++ b/src-java/flowhs-topology/flowhs-storm-topology/src/main/java/org/openkilda/wfm/topology/flowhs/fsm/yflow/update/actions/ValidateYFlowAction.java
@@ -127,10 +127,7 @@ public class ValidateYFlowAction extends
                     format("Unable to map provided sub-flows set onto existing y-flow %s", yFlowId));
         }
 
-        YSubFlow subFlow = yFlow.getSubFlows().stream().findAny()
-                .orElseThrow(() -> new FlowProcessingException(ErrorType.DATA_INVALID,
-                        format("No sub-flows of the y-flow %s were found", yFlowId)));
-        stateMachine.setMainAffinityFlowId(subFlow.getFlow().getAffinityGroupId());
+        stateMachine.setMainAffinityFlowId(getMainAffinityFlowId(yFlow));
 
         List<FlowEndpoint> subFlowEndpoints = targetFlow.getSubFlows().stream()
                 .map(SubFlowDto::getEndpoint)
@@ -145,6 +142,24 @@ public class ValidateYFlowAction extends
         return Optional.empty();
     }
 
+    private String getMainAffinityFlowId(YFlow yFlow) {
+        // TODO: maybe we should add filtering of one switch sub flows into method getSubFlows()
+        YSubFlow multiSwitchFlows = yFlow.getSubFlows().stream()
+                .filter(sub -> !sub.isOneSwitchYFlow(yFlow.getSharedEndpoint().getSwitchId()))
+                .findAny()
+                .orElse(null);
+        if (multiSwitchFlows != null) {
+            return multiSwitchFlows.getFlow().getAffinityGroupId();
+        } else {
+            // if there is no multi switch flows we have to use one switch flow
+            YSubFlow oneSwitchSubFlow = yFlow.getSubFlows().stream()
+                    .findAny()
+                    .orElseThrow(() -> new FlowProcessingException(ErrorType.DATA_INVALID,
+                            format("No sub-flows of the y-flow %s were found", yFlow.getYFlowId())));
+            return oneSwitchSubFlow.getSubFlowId();
+        }
+    }
+
     @Override
     protected String getGenericErrorMessage() {
         return "Could not update y-flow";
diff --git a/src-java/kilda-model/src/main/java/org/openkilda/model/YSubFlow.java b/src-java/kilda-model/src/main/java/org/openkilda/model/YSubFlow.java
index 461140835..53d21e27f 100644
--- a/src-java/kilda-model/src/main/java/org/openkilda/model/YSubFlow.java
+++ b/src-java/kilda-model/src/main/java/org/openkilda/model/YSubFlow.java
@@ -83,6 +83,10 @@ public class YSubFlow implements CompositeDataEntity<YSubFlowData> {
         this.data = builder.build();
     }
 
+    public boolean isOneSwitchYFlow(SwitchId sharedEndpointSwitchId) {
+        return data.getEndpointSwitchId().equals(sharedEndpointSwitchId);
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {