Add tests and expose merge driver map

ImageEngine · Jan 9, 2025 · 9c2a2e1 · 9c2a2e1
1 parent 6104274
commit 9c2a2e1
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 7 deletions.
diff --git a/include/IECoreImage/DisplayDriverServer.h b/include/IECoreImage/DisplayDriverServer.h
@@ -67,6 +67,14 @@ class IECOREIMAGE_API DisplayDriverServer : public IECore::RunTimeTyped
 		using Port = int;
 #endif
 
+		struct MergeDriverInfo
+		{
+			DisplayDriverPtr mergeDriver = nullptr;
+			int mergeCount = 0;
+		};
+
+		using MergeMap = std::map<int, MergeDriverInfo>;
+
 		using PortRange = std::pair<Port, Port>;
 
 		/// A port number of 0 causes a free port to be chosen
@@ -91,6 +99,10 @@ class IECOREIMAGE_API DisplayDriverServer : public IECore::RunTimeTyped
 		static void deregisterPortRange( const std::string &name );
 		static const PortRange &registeredPortRange( const std::string &name );
 
+		// Get the number of merge drivers associated with
+		// the current merge id.
+		static const MergeDriverInfo &getMergeDriverInfo( int mergeId );
+
 	private:
 
 		// Session class

diff --git a/src/IECoreImage/DisplayDriverServer.cpp b/src/IECoreImage/DisplayDriverServer.cpp
@@ -73,13 +73,7 @@ IE_CORE_DEFINERUNTIMETYPED( DisplayDriverServer );
 namespace
 {
 
-struct MergeDriverInfo
-{
-	DisplayDriverPtr mergeDriver = nullptr;
-	int mergeCount = 0;
-};
-
-std::map<int, MergeDriverInfo> g_mergeMap;
+static DisplayDriverServer::MergeMap g_mergeMap;
 
 /* Set the FD_CLOEXEC flag for the given socket descriptor, so that it will not exist on child processes.*/
 static void fixSocketFlags( int socketDesc )
@@ -271,6 +265,17 @@ const DisplayDriverServer::PortRange &DisplayDriverServer::registeredPortRange(
 	return it->second;
 }
 
+const DisplayDriverServer::MergeDriverInfo &DisplayDriverServer::getMergeDriverInfo( int mergeId )
+{
+	auto it = g_mergeMap.find( mergeId );
+	if( it == g_mergeMap.end() )
+	{
+		throw IECore::InvalidArgumentException( "DisplayDriverServer::getMergeCount : " + std::to_string(mergeId) + " is not in the merge map." );
+	}
+
+	return it->second;
+}
+
 DisplayDriverServer::Port DisplayDriverServer::portNumber()
 {
 	return m_data->m_acceptor.local_endpoint().port();

diff --git a/src/IECoreImageBindings/DisplayDriverServerBinding.cpp b/src/IECoreImageBindings/DisplayDriverServerBinding.cpp
@@ -72,6 +72,12 @@ boost::python::tuple registeredPortRange( const std::string &name )
 	return boost::python::make_tuple( range.first, range.second );
 }
 
+boost::python::tuple getMergeDriverInfo( int mergeId )
+{
+	auto mergeDriverInfo = DisplayDriverServer::getMergeDriverInfo( mergeId );
+	return boost::python::make_tuple( mergeDriverInfo.mergeDriver, mergeDriverInfo.mergeCount );
+}
+
 } // namespace
 
 namespace IECoreImageBindings
@@ -89,6 +95,7 @@ void bindDisplayDriverServer()
 		.def( "registerPortRange", &::registerPortRange ).staticmethod( "registerPortRange" )
 		.def( "deregisterPortRange", &DisplayDriverServer::deregisterPortRange ).staticmethod( "deregisterPortRange" )
 		.def( "registeredPortRange", &registeredPortRange ).staticmethod( "registeredPortRange" )
+		.def( "getMergeDriverInfo", &getMergeDriverInfo).staticmethod( "getMergeDriverInfo" )
 	;
 
 }

diff --git a/test/IECoreImage/DisplayDriverServerTest.py b/test/IECoreImage/DisplayDriverServerTest.py
@@ -34,6 +34,7 @@
 
 import unittest
 import sys
+import imath
 
 import IECore
 import IECoreImage
@@ -118,6 +119,47 @@ def testPortRangeRegistry( self ) :
 		s2 = IECoreImage.DisplayDriverServer()
 		self.assertEqual( s2.portNumber(), 45021 )
 
+	def testMergeMap( self ) :
+		server = IECoreImage.DisplayDriverServer( 45001 )
+
+		displayWindow = imath.Box2i( imath.V2i(0,0), imath.V2i(100,100) )
+		dataWindow = imath.Box2i( imath.V2i(0,0), imath.V2i(100,100) )
+
+		params = IECore.CompoundData()
+		params['displayHost'] = IECore.StringData('localhost')
+		params['displayPort'] = IECore.StringData( '45001' )
+		params['displayDriverServer:mergeId'] = IECore.IntData( 42 )
+		params['displayDriverServer:mergeClients'] = IECore.IntData( 2 )
+		params["remoteDisplayType"] = IECore.StringData( "ImageDisplayDriver" )
+
+		idd1 = IECoreImage.ClientDisplayDriver( displayWindow, dataWindow, list( ["R", "G"] ), params )
+		idd2 = IECoreImage.ClientDisplayDriver( displayWindow, dataWindow, list( ["R", "G"] ), params )
+
+		params['displayDriverServer:mergeId'] = IECore.IntData( 666 )
+		params['displayDriverServer:mergeClients'] = IECore.IntData( 1 )
+		idd3 = IECoreImage.ClientDisplayDriver( displayWindow, dataWindow, list( ["R", "G"] ), params )
+
+		# Test that the merge map has merge id 42 and that there are two clients
+		mergeDriverInfo = IECoreImage.DisplayDriverServer.getMergeDriverInfo( 42 )
+		self.assertEqual( mergeDriverInfo[1], 2 )
+
+		# Test that a new merge driver with a new id creates a new display driver
+		mergeDriverInfo2 = IECoreImage.DisplayDriverServer.getMergeDriverInfo( 666 )
+		self.assertNotEqual( mergeDriverInfo[0], mergeDriverInfo2[0] )
+
+		# Test that one image close will remove one client from the merge map
+		idd1.imageClose()
+		mergeDriverInfo = IECoreImage.DisplayDriverServer.getMergeDriverInfo( 42 )
+		self.assertEqual( mergeDriverInfo[1], 1 )
+
+		# Test that after the last image close no clients are left in the map
+		idd2.imageClose()
+		self.assertRaises( RuntimeError, lambda : IECoreImage.DisplayDriverServer.getMergeDriverInfo( 42 ) )
+
+		idd3.imageClose()
+
+		server = None
+
 if __name__ == "__main__":
 	unittest.main()