From b472b5ccb1379ee0917512152c3169a989bc52db Mon Sep 17 00:00:00 2001
From: Andrew Kay <git@ajk.me>
Date: Thu, 4 Nov 2021 18:35:10 -0500
Subject: [PATCH] Allow terminal features to be overridden

---
 oec/terminal.py        | 31 +++++++++++++++++++++++++++-
 tests/test_terminal.py | 47 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/oec/terminal.py b/oec/terminal.py
index ad75312..a4f58aa 100644
--- a/oec/terminal.py
+++ b/oec/terminal.py
@@ -3,6 +3,7 @@ oec.terminal
 ~~~~~~~~~~~~
 """
 
+import os
 import time
 import logging
 from more_itertools import chunked
@@ -171,7 +172,35 @@ def _get_features(interface, device_address):
 
     ids = interface.execute([address_commands(device_address, command) for command in commands])
 
-    return parse_features(ids, commands)
+    features = parse_features(ids, commands)
+
+    # Add override features - for example, this can be used to add an unreported
+    # EAB feature to a IBM 3179 terminal.
+    if 'COAX_FEATURES' in os.environ:
+        for override in os.environ['COAX_FEATURES'].split(','):
+            if '@' not in override:
+                logger.warning(f'Invalid feature override: {override}')
+                continue
+
+            (name, address) = override.split('@')
+
+            try:
+                feature = Feature[name]
+            except KeyError:
+                logger.warning(f'Invalid feature override: {override}')
+                continue
+
+            try:
+                address = int(address)
+            except ValueError:
+                logger.warning(f'Invalid feature override: {override}')
+                continue
+
+            logger.info(f'Adding override feature {feature} @ {address}')
+
+            features[feature] = address
+
+    return features
 
 def _jumbo_write_split_data(data, max_length, first_chunk_max_length_adjustment=-1):
     if max_length is None:
diff --git a/tests/test_terminal.py b/tests/test_terminal.py
index 0c3ca5b..6403c23 100644
--- a/tests/test_terminal.py
+++ b/tests/test_terminal.py
@@ -1,5 +1,5 @@
 import unittest
-from unittest.mock import Mock, create_autospec
+from unittest.mock import Mock, create_autospec, patch
 from coax import Poll, PollAction, TerminalType, Feature, ReadTerminalId, ReadExtendedId, ReadFeatureId
 from coax.protocol import TerminalId
 
@@ -76,7 +76,28 @@ class CreateTerminalTestCase(unittest.TestCase):
 
         self.get_keymap = lambda terminal_id, extended_id: KEYMAP_3278_2
 
-    def test_supported_terminal(self):
+    def test_supported_terminal_with_no_features(self):
+        # Arrange
+        interface = InterfaceWrapper(self.interface)
+
+        self.interface.mock_responses = [
+            (None, ReadTerminalId, None, TerminalId(0b11110100)),
+            (None, ReadExtendedId, None, bytes.fromhex('00 00 00 00'))
+        ]
+
+        # Act
+        terminal = create_terminal(interface, None, None, self.get_keymap)
+
+        # Assert
+        self.assertEqual(terminal.terminal_id.type, TerminalType.CUT)
+        self.assertEqual(terminal.terminal_id.model, 2)
+        self.assertEqual(terminal.terminal_id.keyboard, 15)
+        self.assertEqual(terminal.extended_id, '00000000')
+        self.assertEqual(terminal.display.dimensions, Dimensions(24, 80))
+        self.assertEqual(terminal.features, { })
+        self.assertEqual(terminal.keyboard.keymap.name, '3278-2')
+
+    def test_supported_terminal_with_eab_feature(self):
         # Arrange
         interface = InterfaceWrapper(self.interface)
 
@@ -98,6 +119,28 @@ class CreateTerminalTestCase(unittest.TestCase):
         self.assertEqual(terminal.features, { Feature.EAB: 7 })
         self.assertEqual(terminal.keyboard.keymap.name, '3278-2')
 
+    def test_supported_terminal_features_override(self):
+        # Arrange
+        interface = InterfaceWrapper(self.interface)
+
+        self.interface.mock_responses = [
+            (None, ReadTerminalId, None, TerminalId(0b11110100)),
+            (None, ReadExtendedId, None, bytes.fromhex('00 00 00 00'))
+        ]
+
+        # Act
+        with patch.dict('oec.terminal.os.environ', { 'COAX_FEATURES': 'EAB@7' }):
+            terminal = create_terminal(interface, None, None, self.get_keymap)
+
+        # Assert
+        self.assertEqual(terminal.terminal_id.type, TerminalType.CUT)
+        self.assertEqual(terminal.terminal_id.model, 2)
+        self.assertEqual(terminal.terminal_id.keyboard, 15)
+        self.assertEqual(terminal.extended_id, '00000000')
+        self.assertEqual(terminal.display.dimensions, Dimensions(24, 80))
+        self.assertEqual(terminal.features, { Feature.EAB: 7 })
+        self.assertEqual(terminal.keyboard.keymap.name, '3278-2')
+
     def test_unsupported_terminal_type(self):
         # Arrange
         interface = InterfaceWrapper(self.interface)